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())