## @package batch_lr_loss # Module caffe2.python.layers.batch_lr_loss from caffe2.python import core, schema from caffe2.python.layers.layers import ( ModelLayer, ) from caffe2.python.layers.tags import ( Tags ) import numpy as np class BatchLRLoss(ModelLayer): def __init__( self, model, input_record, name='batch_lr_loss', average_loss=True, jsd_weight=0.0, pos_label_target=1.0, neg_label_target=0.0, homotopy_weighting=False, log_D_trick=False, unjoined_lr_loss=False, uncertainty_penalty=1.0, focal_gamma=0.0, stop_grad_in_focal_factor=False, task_gamma=1.0, task_gamma_lb=0.1, **kwargs ): super(BatchLRLoss, self).__init__(model, name, input_record, **kwargs) self.average_loss = average_loss assert (schema.is_schema_subset( schema.Struct( ('label', schema.Scalar()), ('logit', schema.Scalar()) ), input_record )) self.jsd_fuse = False assert jsd_weight >= 0 and jsd_weight <= 1 if jsd_weight > 0 or homotopy_weighting: assert 'prediction' in input_record self.init_weight(jsd_weight, homotopy_weighting) self.jsd_fuse = True self.homotopy_weighting = homotopy_weighting assert pos_label_target <= 1 and pos_label_target >= 0 assert neg_label_target <= 1 and neg_label_target >= 0 assert pos_label_target >= neg_label_target self.pos_label_target = pos_label_target self.neg_label_target = neg_label_target assert not (log_D_trick and unjoined_lr_loss) self.log_D_trick = log_D_trick self.unjoined_lr_loss = unjoined_lr_loss assert uncertainty_penalty >= 0 self.uncertainty_penalty = uncertainty_penalty self.tags.update([Tags.EXCLUDE_FROM_PREDICTION]) self.output_schema = schema.Scalar( np.float32, self.get_next_blob_reference('output') ) self.focal_gamma = focal_gamma self.stop_grad_in_focal_factor = stop_grad_in_focal_factor self.apply_exp_decay = False if task_gamma < 1.0: self.apply_exp_decay = True self.task_gamma_cur = self.create_param( param_name=('%s_task_gamma_cur' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': 1.0, 'dtype': core.DataType.FLOAT } ), optimizer=self.model.NoOptim, ) self.task_gamma = self.create_param( param_name=('%s_task_gamma' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': task_gamma, 'dtype': core.DataType.FLOAT } ), optimizer=self.model.NoOptim, ) self.task_gamma_lb = self.create_param( param_name=('%s_task_gamma_lb' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': task_gamma_lb, 'dtype': core.DataType.FLOAT } ), optimizer=self.model.NoOptim, ) def init_weight(self, jsd_weight, homotopy_weighting): if homotopy_weighting: self.mutex = self.create_param( param_name=('%s_mutex' % self.name), shape=None, initializer=('CreateMutex', ), optimizer=self.model.NoOptim, ) self.counter = self.create_param( param_name=('%s_counter' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': 0, 'dtype': core.DataType.INT64 } ), optimizer=self.model.NoOptim, ) self.xent_weight = self.create_param( param_name=('%s_xent_weight' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': 1., 'dtype': core.DataType.FLOAT } ), optimizer=self.model.NoOptim, ) self.jsd_weight = self.create_param( param_name=('%s_jsd_weight' % self.name), shape=[1], initializer=( 'ConstantFill', { 'value': 0., 'dtype': core.DataType.FLOAT } ), optimizer=self.model.NoOptim, ) else: self.jsd_weight = self.model.add_global_constant( '%s_jsd_weight' % self.name, jsd_weight ) self.xent_weight = self.model.add_global_constant( '%s_xent_weight' % self.name, 1. - jsd_weight ) def update_weight(self, net): net.AtomicIter([self.mutex, self.counter], [self.counter]) # iter = 0: lr = 1; # iter = 1e6; lr = 0.5^0.1 = 0.93 # iter = 1e9; lr = 1e-3^0.1 = 0.50 net.LearningRate([self.counter], [self.xent_weight], base_lr=1.0, policy='inv', gamma=1e-6, power=0.1,) net.Sub( [self.model.global_constants['ONE'], self.xent_weight], [self.jsd_weight] ) return self.xent_weight, self.jsd_weight def add_ops(self, net): # numerically stable log-softmax with crossentropy label = self.input_record.label() # mandatory cast to float32 # self.input_record.label.field_type().base is np.float32 but # label type is actually int label = net.Cast( label, net.NextScopedBlob('label_float32'), to=core.DataType.FLOAT) label = net.ExpandDims(label, net.NextScopedBlob('expanded_label'), dims=[1]) if self.pos_label_target != 1.0 or self.neg_label_target != 0.0: label = net.StumpFunc( label, net.NextScopedBlob('smoothed_label'), threshold=0.5, low_value=self.neg_label_target, high_value=self.pos_label_target, ) xent = net.SigmoidCrossEntropyWithLogits( [self.input_record.logit(), label], net.NextScopedBlob('cross_entropy'), log_D_trick=self.log_D_trick, unjoined_lr_loss=self.unjoined_lr_loss ) if self.focal_gamma != 0: label = net.StopGradient( [label], [net.NextScopedBlob('label_stop_gradient')], ) prediction = self.input_record.prediction() # focal loss = (y(1-p) + p(1-y))^gamma * original LR loss # y(1-p) + p(1-y) = y + p - 2 * yp y_plus_p = net.Add( [prediction, label], net.NextScopedBlob("y_plus_p"), ) yp = net.Mul([prediction, label], net.NextScopedBlob("yp")) two_yp = net.Scale(yp, net.NextScopedBlob("two_yp"), scale=2.0) y_plus_p_sub_two_yp = net.Sub( [y_plus_p, two_yp], net.NextScopedBlob("y_plus_p_sub_two_yp") ) focal_factor = net.Pow( y_plus_p_sub_two_yp, net.NextScopedBlob("y_plus_p_sub_two_yp_power"), exponent=float(self.focal_gamma), ) if self.stop_grad_in_focal_factor is True: focal_factor = net.StopGradient( [focal_factor], [net.NextScopedBlob("focal_factor_stop_gradient")], ) xent = net.Mul( [xent, focal_factor], net.NextScopedBlob("focallossxent") ) if self.apply_exp_decay: net.Mul( [self.task_gamma_cur, self.task_gamma], self.task_gamma_cur ) task_gamma_multiplier = net.Max( [self.task_gamma_cur, self.task_gamma_lb], net.NextScopedBlob("task_gamma_cur_multiplier") ) xent = net.Mul( [xent, task_gamma_multiplier], net.NextScopedBlob("expdecayxent") ) # fuse with JSD if self.jsd_fuse: jsd = net.BernoulliJSD( [self.input_record.prediction(), label], net.NextScopedBlob('jsd'), ) if self.homotopy_weighting: self.update_weight(net) loss = net.WeightedSum( [xent, self.xent_weight, jsd, self.jsd_weight], net.NextScopedBlob('loss'), ) else: loss = xent if 'log_variance' in self.input_record.fields: # mean (0.5 * exp(-s) * loss + 0.5 * penalty * s) log_variance_blob = self.input_record.log_variance() log_variance_blob = net.ExpandDims( log_variance_blob, net.NextScopedBlob('expanded_log_variance'), dims=[1] ) neg_log_variance_blob = net.Negative( [log_variance_blob], net.NextScopedBlob('neg_log_variance') ) # enforce less than 88 to avoid OverflowError neg_log_variance_blob = net.Clip( [neg_log_variance_blob], net.NextScopedBlob('clipped_neg_log_variance'), max=88.0 ) exp_neg_log_variance_blob = net.Exp( [neg_log_variance_blob], net.NextScopedBlob('exp_neg_log_variance') ) exp_neg_log_variance_loss_blob = net.Mul( [exp_neg_log_variance_blob, loss], net.NextScopedBlob('exp_neg_log_variance_loss') ) penalized_uncertainty = net.Scale( log_variance_blob, net.NextScopedBlob("penalized_unceratinty"), scale=float(self.uncertainty_penalty) ) loss_2x = net.Add( [exp_neg_log_variance_loss_blob, penalized_uncertainty], net.NextScopedBlob('loss') ) loss = net.Scale(loss_2x, net.NextScopedBlob("loss"), scale=0.5) if 'weight' in self.input_record.fields: weight_blob = self.input_record.weight() if self.input_record.weight.field_type().base != np.float32: weight_blob = net.Cast( weight_blob, weight_blob + '_float32', to=core.DataType.FLOAT ) weight_blob = net.StopGradient( [weight_blob], [net.NextScopedBlob('weight_stop_gradient')], ) loss = net.Mul( [loss, weight_blob], net.NextScopedBlob('weighted_cross_entropy'), ) if self.average_loss: net.AveragedLoss(loss, self.output_schema.field_blobs()) else: net.ReduceFrontSum(loss, self.output_schema.field_blobs())