| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970 |
- 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()
- # Load keys into indices
- 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)
|
