mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 13:44:15 +08:00
ecace9eb21
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/12200 moved crf_viterbi_op, copied crf_predict and crf_viterbi_test to oss Reviewed By: Yangqing Differential Revision: D10118341 fbshipit-source-id: 51e30e57d280d6ca75fc0b488f743794f23b589f
34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
from __future__ import absolute_import, division, print_function, unicode_literals
|
|
|
|
import numpy as np
|
|
from caffe2.python.crf import CRFWithLoss
|
|
|
|
|
|
def crf_update_predictions(model, crf_with_loss, classes):
|
|
return apply_crf(
|
|
model.param_init_net,
|
|
model.net,
|
|
crf_with_loss.transitions,
|
|
classes,
|
|
crf_with_loss.num_classes,
|
|
)
|
|
|
|
|
|
def apply_crf(init_net, net, transitions, predictions, num_classes):
|
|
padded_classes = CRFWithLoss.pad_predictions(
|
|
predictions, init_net, net, num_classes
|
|
)
|
|
bestPath = net.ViterbiPath([padded_classes, transitions])
|
|
new_padded_classes = net.SwapBestPath([padded_classes, bestPath])
|
|
# Revert the effect of pad_predictions by removing the last two rows and
|
|
# the last two columns
|
|
new_classes = net.RemovePadding(
|
|
[new_padded_classes], padding_width=1, end_padding_width=1
|
|
)
|
|
slice_starts = np.array([0, 0]).astype(np.int32)
|
|
slice_ends = np.array([-1, -3]).astype(np.int32)
|
|
slice_starts = net.GivenTensorIntFill([], shape=[2], values=slice_starts)
|
|
slice_ends = net.GivenTensorIntFill([], shape=[2], values=slice_ends)
|
|
new_classes = net.Slice([new_classes, slice_starts, slice_ends])
|
|
return new_classes
|