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