|
|
|
|
|
|
|
|
|
|
| import numpy as np |
|
|
| from caffe2.python import core, schema |
| from caffe2.python.layers.layers import ModelLayer |
|
|
|
|
| class MapToRange(ModelLayer): |
| """ |
| This layer aims to build a mapping from raw keys to indices within [0, max_index). |
| The mapping is continuously built during training. The mapping will be frozen during |
| evaluation and prediction. Unseen keys will be assigned to index 0. |
| """ |
|
|
| def __init__( |
| self, model, |
| input_record, |
| max_index, |
| name='map_to_range', |
| **kwargs |
| ): |
| super(MapToRange, self).__init__(model, name, input_record, **kwargs) |
|
|
| assert max_index > 0 |
| assert isinstance(input_record, schema.Scalar) |
|
|
| self.max_index = max_index |
|
|
| self.handler = self.create_param( |
| param_name='handler', |
| shape=[], |
| initializer=('LongIndexCreate', {'max_elements': self.max_index}), |
| optimizer=model.NoOptim |
| ) |
|
|
| self.output_schema = schema.Struct( |
| ('indices', schema.Scalar( |
| np.int64, self.get_next_blob_reference("indices") |
| )), |
| ('handler', schema.Scalar( |
| np.void, self.handler |
| )), |
| ) |
|
|
| def add_train_ops(self, net): |
| if self.input_record.field_type().base != np.int64: |
| keys = net.Cast( |
| self.input_record(), |
| net.NextScopedBlob("indices_before_mapping"), |
| to=core.DataType.INT64 |
| ) |
| else: |
| keys = self.input_record() |
|
|
| |
| indices = net.IndexGet([self.handler, keys], |
| self.output_schema.indices()) |
|
|
| net.StopGradient(indices, indices) |
|
|
| def add_eval_ops(self, net): |
| net.IndexFreeze(self.handler, self.handler) |
| self.add_train_ops(net) |
|
|
| def add_ops(self, net): |
| self.add_eval_ops(net) |
|
|
