| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | import logging |
| | import numpy as np |
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ( |
| | get_categorical_limit, |
| | ModelLayer, |
| | ) |
| |
|
| | from caffe2.python.layers.tags import Tags |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class PositionWeighted(ModelLayer): |
| | def __init__(self, model, input_record, weight_optim=None, |
| | name="position_weights"): |
| | super(PositionWeighted, self).__init__(model, name, input_record) |
| |
|
| | assert isinstance(input_record, schema.List), "Incorrect input type" |
| | length_metadata = input_record.lengths.metadata |
| | max_length = (length_metadata.categorical_limit if length_metadata is |
| | not None else None) |
| | if max_length is not None: |
| | self.shape = max_length |
| | else: |
| | self.shape = get_categorical_limit(input_record) |
| | logger.warning( |
| | '{}: categorical_limit of lengths is not available, using ' |
| | 'categorical_limit of the keys: {}'.format( |
| | str(input_record.lengths()), self.shape)) |
| |
|
| | self.pos_w = self.create_param(param_name='pos_w', |
| | shape=[self.shape, ], |
| | initializer=('ConstantFill', {'value': 1.0}), |
| | optimizer=weight_optim) |
| |
|
| | self.output_schema = schema.Struct( |
| | ('position_weights', |
| | schema.Scalar((np.float32, self.shape), |
| | self.get_next_blob_reference("pos_w_gather"))) |
| | ) |
| |
|
| | self.tags.update({Tags.HANDLE_AS_SPARSE_LAYER}) |
| |
|
| | def get_memory_usage(self): |
| | return self.shape |
| |
|
| | def add_ops(self, net): |
| | inc_seq = net.LengthsRangeFill( |
| | [self.input_record.lengths()], |
| | self.input_record.lengths() + '_pos_w_seq' |
| | ) |
| |
|
| | net.Gather( |
| | [self.pos_w, inc_seq], |
| | self.output_schema.position_weights.field_blobs()) |
| |
|
