| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema, core |
| | from caffe2.python.layers.layers import ( |
| | ModelLayer, |
| | ) |
| | from caffe2.python.layers.tags import ( |
| | Tags |
| | ) |
| | import numpy as np |
| |
|
| |
|
| | class MarginRankLoss(ModelLayer): |
| |
|
| | def __init__(self, model, input_record, name='margin_rank_loss', |
| | margin=0.1, average_loss=False, **kwargs): |
| | super(MarginRankLoss, self).__init__(model, name, input_record, **kwargs) |
| | assert margin >= 0, ('For hinge loss, margin should be no less than 0') |
| | self._margin = margin |
| | self._average_loss = average_loss |
| | assert schema.is_schema_subset( |
| | schema.Struct( |
| | ('pos_prediction', schema.Scalar()), |
| | ('neg_prediction', schema.List(np.float32)), |
| | ), |
| | input_record |
| | ) |
| | self.tags.update([Tags.EXCLUDE_FROM_PREDICTION]) |
| | self.output_schema = schema.Scalar( |
| | np.float32, |
| | self.get_next_blob_reference('output')) |
| |
|
| | def add_ops(self, net): |
| | neg_score = self.input_record.neg_prediction['values']() |
| |
|
| | pos_score = net.LengthsTile( |
| | [ |
| | self.input_record.pos_prediction(), |
| | self.input_record.neg_prediction['lengths']() |
| | ], |
| | net.NextScopedBlob('pos_score_repeated') |
| | ) |
| | const_1 = net.ConstantFill( |
| | neg_score, |
| | net.NextScopedBlob('const_1'), |
| | value=1, |
| | dtype=core.DataType.INT32 |
| | ) |
| | rank_loss = net.MarginRankingCriterion( |
| | [pos_score, neg_score, const_1], |
| | net.NextScopedBlob('rank_loss'), |
| | margin=self._margin, |
| | ) |
| | if self._average_loss: |
| | net.AveragedLoss(rank_loss, self.output_schema.field_blobs()) |
| | else: |
| | net.ReduceFrontSum(rank_loss, self.output_schema.field_blobs()) |
| |
|
