/usr/local/lib64/python3.6/site-packages/caffe2/python
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/crf.py (13250B)
## @package crf
# Module caffe2.python.crf
import numpy as np
from caffe2.python import brew, core, model_helper, recurrent
"""
Due to a limitation in ReccurentNetworkOp, this layer only supports batch_size=1
In order to support batch_size > 1, we will have to implement the CRFUnit
and its gradient in C++ and handle the different batches there.
"""
class CRFWithLoss(object):
def __init__(self, model, num_classes, transitions_blob=None):
self.model = model
self.num_classes = num_classes
self.num_classes_padded = num_classes + 2 # After adding BOS and EOS
if not transitions_blob:
transitions_blob = self.model.param_init_net.UniformFill(
[],
[core.ScopedBlobReference("crf_transitions")],
shape=[self.num_classes_padded, self.num_classes_padded],
min=-1.0,
max=1.0,
)
self.transitions = transitions_blob
self.model.params.append(self.transitions)
def crf_loss(self, predictions, labels, seq_lengths=None):
# Since the transitions matrix is a shared parameter, need to
# take a snapshot of it at the beginning since it can be updated
# in between the operators that uses it when doing parallel updates
transitions_snapshot = self.model.net.Copy(
self.transitions, core.ScopedBlobReference("transitions_snapshot")
)
# Compute best path unary score from the logits
path_unary_score = self._gather_entries_sum(
predictions, labels, self.num_classes
)
# Append BOS and EOS entries to the predictions and labels
predictions = CRFWithLoss.pad_predictions(
predictions, self.model.param_init_net, self.model.net, self.num_classes
)
labels = CRFWithLoss.pad_labels(
labels, self.model.param_init_net, self.model.net, self.num_classes
)
# Compute best path binary scores from the transitions matrix
path_binary_score = self._path_binary_scores(
labels, transitions_snapshot, seq_lengths
)
path_total_score = self.model.net.Add(
[path_binary_score, path_unary_score],
core.ScopedBlobReference("path_total"),
)
# Compute all paths score
zero_index = self.model.param_init_net.ConstantFill([], shape=[1], value=0)
initial_state = self.model.net.Gather(
[predictions, zero_index],
core.ScopedBlobReference("rnn_initial"),
dense_gradient=True,
)
input_data, _ = self.model.net.RemovePadding(
[predictions], padding_width=1, end_padding_width=0, outputs=2
)
input_data = self.model.net.ExpandDims(
[input_data], core.ScopedBlobReference("rnn_input_data"), dims=[1]
)
# Due to a bug in RecurrentNetworkGradientOp, we need to copy the
# transitions blob before sending it to the recurrent network
transitions_copy = self.model.net.Copy(
transitions_snapshot, core.ScopedBlobReference("transitions_copy")
)
all_paths_scores = self._crf_forward(
input_data, initial_state, transitions_copy
)
loss = self.model.net.Sub(
[all_paths_scores, path_total_score], core.ScopedBlobReference("crf_loss")
)
return loss
def _path_binary_scores(self, labels, transitions, seq_lengths=None):
column_ids, _ = self.model.net.RemovePadding(
[labels], outputs=2, padding_width=1, end_padding_width=0
)
row_ids, _ = self.model.net.RemovePadding(
[labels], outputs=2, padding_width=0, end_padding_width=1
)
# Since there is no multi-dimensional gather, I flatten the matrix to
# a 1-d vector and transform the ids to (row_ids * num_columns +
# column_ids) and do gather in 1-d
num_columns_blob = self.model.net.ConstantFill(
[row_ids], value=self.num_classes_padded
)
flattened_ids = self.model.net.Mul([row_ids, num_columns_blob])
flattened_ids = self.model.net.Add([flattened_ids, column_ids])
flattened_transitions = self.model.net.FlattenToVec([transitions])
entries = self.model.net.Gather(
[flattened_transitions, flattened_ids], dense_gradient=True
)
return self.model.ReduceFrontSum(entries)
def _gather_entries_sum(self, in_data, indices, index_size):
indices = self.model.net.Cast([indices], to="int64")
index_size_blob = self.model.param_init_net.ConstantFill(
[], shape=[1], value=index_size
)
query_one_hot = self.model.net.OneHot([indices, index_size_blob])
flattend_query = self.model.net.FlattenToVec(query_one_hot)
flattend_data = self.model.net.FlattenToVec(in_data)
query_scores = self.model.net.DotProduct([flattend_query, flattend_data])
final_sum = self.model.net.ReduceFrontSum([query_scores])
return final_sum
def _crf_forward(
self, input_blob, initial_state, transitions_copy, seq_lengths=None
):
# Build the RNN net and get the last timestep output
out_last = self.build_crf_net(input_blob, initial_state, transitions_copy)
out_last, _ = self.model.net.Reshape(
[out_last], outputs=2, shape=(self.num_classes_padded,)
)
zero_segment_id = self.model.param_init_net.ConstantFill(
[], value=0, shape=[self.num_classes_padded], dtype=core.DataType.INT32
)
# Compute the accumulated total score of all the paths
accum_score = self.model.net.SortedSegmentRangeLogSumExp(
[out_last, zero_segment_id]
)
accum_score, _ = self.model.net.Reshape(accum_score, outputs=2, shape=())
return accum_score
def build_crf_net(self, input_blob, initial_state, transitions):
"""
Adds the crf_net recurrent operator to the model.
model: model_helper.ModelHelper object new operators would be added
to
input_blob: the input sequence in a format T x N x D
where T is sequence size, N - batch size and D - input dimension
##Only supports batch-size 1##
seq_lengths: blob containing sequence lengths (unused)
"""
scope = "crf_net"
def s(name):
""
# We have to manually scope due to our internal/external blob
# relationships.
return "{}/{}".format(str(scope), str(name))
step_model = model_helper.ModelHelper(name="crf_step", param_model=self.model)
input_t, cell_t_prev, _ = step_model.net.AddExternalInputs(
core.ScopedBlobReference("input_t"),
core.ScopedBlobReference("cell_t_prev"),
transitions,
)
zero_segment_id = step_model.param_init_net.ConstantFill(
[],
[s("zero_segment_id")],
value=0,
shape=[self.num_classes_padded],
dtype=core.DataType.INT32,
)
# A hack to bypass model cloning for test
step_model.param_init_net.AddExternalOutput(zero_segment_id)
""" the CRF step """
# Do tile
prev_transpose = brew.transpose(
step_model, cell_t_prev, [s("prev_transpose")], axes=(0, 2, 1)
)
prev_tiled = step_model.net.Tile(
prev_transpose, [s("prev_tiled")], tiles=self.num_classes_padded, axis=2
)
input_t_tiled = step_model.net.Tile(
input_t, [s("input_t_tiled")], tiles=self.num_classes_padded, axis=1
)
input_with_prev = step_model.net.Add(
[prev_tiled, input_t_tiled], [s("input_with_prev")]
)
all_with_transitions = step_model.net.Add(
[input_with_prev, transitions],
[s("prev_with_transitions")],
broadcast=1,
use_grad_hack=1,
)
all_with_transitions_reshaped, _ = step_model.net.Reshape(
all_with_transitions,
[s("all_with_transitions_reshaped"), s("all_with_transitions_orig")],
shape=(self.num_classes_padded, self.num_classes_padded),
)
cell_t = step_model.net.SortedSegmentRangeLogSumExp(
[all_with_transitions_reshaped, zero_segment_id], [s("cell_t")]
)
step_model.net.AddExternalOutputs(cell_t)
""" recurrent network """
cell_input_blob = initial_state
out_all, out_last = recurrent.recurrent_net(
net=self.model.net,
cell_net=step_model.net,
inputs=[(input_t, input_blob)],
initial_cell_inputs=[(cell_t_prev, cell_input_blob)],
links={cell_t_prev: cell_t},
scope=scope,
outputs_with_grads=(1,),
)
return out_last
def update_predictions(self, classes):
def crf_update_predictions_op(inputs, outputs):
# This operator will compute the best path of classes by performing
# Viterbi decoding and then updates the predictions to make the tag
# On the best path has the highest score among the others
predictions = inputs[0].data
transitions = inputs[1].data
predictions = inputs[0].data
predictions_shape = inputs[0].shape
outputs[0].reshape(predictions_shape)
trellis = np.zeros(predictions_shape)
backpointers = np.zeros(predictions_shape, dtype=np.int32)
trellis[0] = predictions[0]
for t in range(1, predictions_shape[0]):
v = np.expand_dims(trellis[t - 1], 1) + transitions
trellis[t] = predictions[t] + np.max(v, 0)
backpointers[t] = np.argmax(v, 0)
viterbi = [np.argmax(trellis[-1])]
for bp in reversed(backpointers[1:]):
viterbi.append(bp[viterbi[-1]])
viterbi.reverse()
new_predictions = np.zeros(predictions_shape)
old_bests = []
for i, w_predictions in enumerate(predictions):
# Get the current tag with the maximum score
new_predictions[i] = predictions[i]
old_best = np.argmax(w_predictions)
old_bests.append(old_best)
# Swap the scores of the current best tag and the tag on the
# Viterbi path
w_predictions[viterbi[i]], w_predictions[old_best] = (
w_predictions[old_best],
w_predictions[viterbi[i]],
)
new_predictions[i] = w_predictions
# Remove the BOS and EOS entries from the predictions matrix
orig_predictions = new_predictions[1:-1, 0:-2]
outputs[0].reshape(orig_predictions.shape)
outputs[0].data[...] = orig_predictions
padded_classes = CRFWithLoss.pad_predictions(
classes, self.model.param_init_net, self.model.net, self.num_classes
)
new_classes = self.model.net.Python(crf_update_predictions_op)(
[padded_classes, self.transitions],
core.ScopedBlobReference("post_crf_classes"),
)
return new_classes
@staticmethod
def pad_labels(labels, init_net, net, num_classes):
bos_i = num_classes
eos_i = num_classes + 1
bos_i_b = init_net.ConstantFill([], shape=[1], value=bos_i)
eos_i_b = init_net.ConstantFill([], shape=[1], value=eos_i)
labels = net.Cast([labels], to="int64")
padded_labels, _ = net.Concat([bos_i_b, labels, eos_i_b], axis=0, outputs=2)
return padded_labels
@staticmethod
def pad_predictions(predictions, init_net, net, num_classes):
# This function will introduce two labels for beginning of sequence
# And end of sequence, it will make the necessary udpates to the
# the predictions blob
low_score = -1000.0 # An arbitray very low number
b_scores = np.array([[low_score] * num_classes + [0, low_score]]).astype(
np.float32
)
e_scores = np.array([[low_score] * num_classes + [low_score, 0]]).astype(
np.float32
)
b_scores = init_net.GivenTensorFill(
[], "b_scores", shape=[1, num_classes + 2], values=b_scores
)
e_scores = init_net.GivenTensorFill(
[], "e_scores", shape=[1, num_classes + 2], values=e_scores
)
zero_index = net.ConstantFill([], shape=[1], value=0)
length = net.Gather([net.Shape([predictions]), zero_index])
length = net.Cast(length, to="int32")
t_range = net.LengthsRangeFill(length)
padding = net.ConstantFill([t_range], value=low_score)
padding = net.ExpandDims(padding, dims=[1])
padded_predictions, _ = net.Concat(
[predictions, padding, padding], outputs=2, axis=1
)
padded_predictions_concat, _ = net.Concat(
[b_scores, padded_predictions, e_scores], outputs=2, axis=0
)
return padded_predictions_concat