def ExtractPredictorNet(model, inputs, outputs, device):
'''
Returns (net, params) that can be exported to be used as a prediction
net.
'''
master_device = model._devices[0]
prefix = "gpu_{}/".format(master_device)
prefix_inputs = [prefix + str(b) for b in inputs]
prefix_outputs = [prefix + str(b) for b in outputs]
(predictor_net, export_blobs) = model_helper.ExtractPredictorNet(
net_proto=model.net.Proto(),
input_blobs=prefix_inputs,
output_blobs=prefix_outputs,
device=device,
renames={
a: b
for (a, b) in zip(prefix_inputs + prefix_outputs, inputs + outputs)
},
)
return (predictor_net, export_blobs)
def build_net(
self,
base_learning_rate=0.1 # base_learning_rate * seq_size
):
log.debug('>>> Building Mask-RNN')
model = model_helper.ModelHelper(name="mask_rnn")
hidden_init = model.net.AddExternalInputs('hidden_init', )
# TODO: do I still need this?
model.net.AddExternalInputs(
'input_blob',
'seq_lengths',
'target',
)
# Add external inputs (read directly from the database)
# the dimension of class_target_mask: [BATCH_SIZE, SEQ_LEN, 1]
# the dimension of regre_target_mask: [BATCH_SIZE, SEQ_LEN, regre_output_dim]
(seq_lengths, _input_blob, _class_target, _regre_target,
_class_target_mask, _regre_target_mask) = build_input_reader(
model,
self.db_name,
'minidb', [
'seq_lengths', 'input_blob_batch_first',
'class_target_batch_first', 'regre_target_batch_first',
'class_target_mask_batch_first',
'regre_target_mask_batch_first'
],
batch_size=self.batch_size,
data_type='train')
# In order to put into batches, the input_blob is
# [BATCH_SIZE, SEQ_LEN, INPUT_DIM]
# i.e. the first dim is the batch size
# However the required input dim is:
# [SEQ_LEN, BATCH_SIZE, INPUT_DIM]
input_blob = model.net.Transpose([_input_blob],
'input_blob',
axes=[1, 0, 2])
class_target = model.net.Transpose([_class_target],
'class_target',
axes=[1, 0, 2])
regre_target = model.net.Transpose([_regre_target],
'regre_target',
axes=[1, 0, 2])
class_target_mask = model.net.Transpose([_class_target_mask],
'class_target_mask',
axes=[1, 0, 2])
regre_target_mask = model.net.Transpose([_regre_target_mask],
'regre_target_mask',
axes=[1, 0, 2])
hidden_output_all, self.hidden_output = MaskGRU(model,
input_blob,
seq_lengths,
(hidden_init, ),
self.input_dim,
self.hidden_size,
scope="MaskRNN")
# axis is 2 as first two are T (time) and N (batch size)
# multi-task learning: regression
regre_output = brew.fc(model,
hidden_output_all,
None,
dim_in=self.hidden_size,
dim_out=self.regre_output_dim,
axis=2)
# multi-task learning: classification
class_output = brew.fc(model,
hidden_output_all,
None,
dim_in=self.hidden_size,
dim_out=self.class_output_dim,
axis=2)
# softmax head for testing only
class_softmax_output = model.net.Softmax(class_output,
'class_softmax_output',
axis=2)
# Get the predict net
(self.net_store['predict'],
self.external_inputs) = model_helper.ExtractPredictorNet(
model.net.Proto(),
[input_blob, seq_lengths, hidden_init],
[class_softmax_output, regre_output],
)
# Then, we add loss and gradient ops
# We treat them as one big batch of size T * N
# we use the logit of classification head
# class_output_reshaped, _ = model.net.Reshape(
# class_output, ['class_output_reshaped', '_class_output_shape'],
# shape=[-1, self.class_output_dim])
class_softmax_output_reshaped, _ = model.net.Reshape(
class_softmax_output,
['class_softmax_output_reshaped', '_class_output_shape'],
shape=[-1, self.class_output_dim])
regre_output_reshaped, _ = model.net.Reshape(
regre_output, ['regre_output_reshaped', '_regre_output_shape'],
shape=[-1, self.regre_output_dim])
class_target_reshaped, _ = model.net.Reshape(
class_target, ['class_target_reshaped', '_class_target_shape'],
shape=[-1, self.class_output_dim])
regre_target_reshaped, _ = model.net.Reshape(
regre_target, ['regre_target_reshaped', '_regre_target_shape'],
shape=[-1, self.regre_output_dim])
class_target_mask_reshaped, _ = model.net.Reshape(
class_target_mask,
['class_target_mask_reshaped', '_class_target_mask_shape'],
shape=[-1, 1])
regre_target_mask_reshaped, _ = model.net.Reshape(
regre_target_mask,
['regre_target_mask_reshaped', '_regre_target_mask_shape'],
shape=[-1, self.regre_output_dim])
# stop gradient to label and mask
class_target_reshaped = model.net.StopGradient(
class_target_reshaped, 'stopped_class_target_reshaped')
regre_target_reshaped = model.net.StopGradient(
regre_target_reshaped, 'stopped_regre_target_reshaped')
class_target_mask_reshaped = model.net.StopGradient(
class_target_mask_reshaped, 'stopped_class_target_mask_reshaped')
regre_target_mask_reshaped = model.net.StopGradient(
regre_target_mask_reshaped, 'stopped_regre_target_mask_reshaped')
# model.net.Print([class_output_reshaped], 'print', to_file=0)
# classification error
# combined softmax and log likelihood for numerical stability
# weighted by class_target_mask_reshaped
#
# _, class_average_loss = model.net.SoftmaxWithLoss(
# [class_output_reshaped, class_target_reshaped, class_target_mask_reshaped],
# ['_train_softmax_ouput', 'class_average_loss'], label_prob=1
# )
#
class_l2_dist = model.net.SquaredL2Distance(
[class_softmax_output_reshaped, class_target_reshaped],
'class_l2_dist')
class_target_mask_reshaped = model.net.Squeeze(
class_target_mask_reshaped, 'squeezed_class_target_mask', dims=[1])
masked_class_l2_dist = model.net.Mul(
[class_target_mask_reshaped, class_l2_dist],
'masked_class_l2_dist')
class_average_loss = model.net.AveragedLoss(masked_class_l2_dist,
'class_average_loss')
# regression error
# mask need to be applied to *each* individual dimension of output vector
regre_output_reshaped_list = model.net.Split(
[regre_output_reshaped],
[
'regre_output_reshaped_' + str(i)
for i in range(self.regre_output_dim)
],
axis=1, # has been reshaped to 2D tensor
)
regre_target_reshaped_list = model.net.Split(
[regre_target_reshaped],
[
'regre_target_reshaped_' + str(i)
for i in range(self.regre_output_dim)
],
axis=1, # has been reshaped to 2D tensor
)
regre_target_mask_reshaped_list = model.net.Split(
[regre_target_mask_reshaped],
[
'regre_target_mask_reshaped_' + str(i)
for i in range(self.regre_output_dim)
],
axis=1, # has been reshaped to 2D tensor
)
regre_average_loss_lst = []
i = 0
for o, t, m in zip(regre_output_reshaped_list,
regre_target_reshaped_list,
regre_target_mask_reshaped_list):
l2_dist = model.net.SquaredL2Distance([o, t], 'l2_dist_' + str(i))
m = model.net.Squeeze(m,
'squeezed_regre_target_mask_' + str(i),
dims=[1])
masked_l2_dist = model.net.Mul([m, l2_dist],
'masked_l2_dist_' + str(i))
# masked_l2_dist = l2_dist
regre_average_loss_lst.append(
model.net.AveragedLoss(masked_l2_dist,
'regre_average_loss_' + str(i)))
i += 1
assert i == self.regre_output_dim, 'output dim != # of loss split'
# Training net
model.AddGradientOperators([class_average_loss] +
regre_average_loss_lst)
build_adam(
model,
base_learning_rate=base_learning_rate * self.seq_size,
)
self.model = model
self.predictions = [class_softmax_output, regre_output]
self.loss = [class_average_loss] + regre_average_loss_lst
for loss in self.loss:
loss = str(loss)
self.reports[loss] = []
# Create a net to copy hidden_output to hidden_init
prepare_state = core.Net("prepare_state")
prepare_state.Copy(self.hidden_output, hidden_init)
self.net_store['prepare'] = prepare_state
self.net_store['train'] = core.Net(model.net.Proto())
