| |
| |
|
|
|
|
|
|
|
|
|
|
| from caffe2.python import core, schema |
| from caffe2.python.layers.layers import ModelLayer |
|
|
|
|
| class ReservoirSampling(ModelLayer): |
| """ |
| Collect samples from input record w/ reservoir sampling. If you have complex |
| data, use PackRecords to pack it before using this layer. |
| |
| This layer is not thread safe. |
| """ |
|
|
| def __init__(self, model, input_record, num_to_collect, |
| name='reservoir_sampling', **kwargs): |
| super(ReservoirSampling, self).__init__( |
| model, name, input_record, **kwargs) |
| assert num_to_collect > 0 |
| self.num_to_collect = num_to_collect |
|
|
| self.reservoir = self.create_param( |
| param_name='reservoir', |
| shape=[0], |
| initializer=('ConstantFill',), |
| optimizer=model.NoOptim, |
| ) |
| self.num_visited_blob = self.create_param( |
| param_name='num_visited', |
| shape=[], |
| initializer=('ConstantFill', { |
| 'value': 0, |
| 'dtype': core.DataType.INT64, |
| }), |
| optimizer=model.NoOptim, |
| ) |
| self.mutex = self.create_param( |
| param_name='mutex', |
| shape=[], |
| initializer=('CreateMutex',), |
| optimizer=model.NoOptim, |
| ) |
|
|
| self.extra_input_blobs = [] |
| self.extra_output_blobs = [] |
| if 'object_id' in input_record: |
| object_to_pos = self.create_param( |
| param_name='object_to_pos', |
| shape=None, |
| initializer=('CreateMap', { |
| 'key_dtype': core.DataType.INT64, |
| 'valued_dtype': core.DataType.INT32, |
| }), |
| optimizer=model.NoOptim, |
| ) |
| pos_to_object = self.create_param( |
| param_name='pos_to_object', |
| shape=[0], |
| initializer=('ConstantFill', { |
| 'value': 0, |
| 'dtype': core.DataType.INT64, |
| }), |
| optimizer=model.NoOptim, |
| ) |
| self.extra_input_blobs.append(input_record.object_id()) |
| self.extra_input_blobs.extend([object_to_pos, pos_to_object]) |
| self.extra_output_blobs.extend([object_to_pos, pos_to_object]) |
|
|
| self.output_schema = schema.Struct( |
| ( |
| 'reservoir', |
| schema.from_blob_list(input_record.data, [self.reservoir]) |
| ), |
| ('num_visited', schema.Scalar(blob=self.num_visited_blob)), |
| ('mutex', schema.Scalar(blob=self.mutex)), |
| ) |
|
|
| def add_ops(self, net): |
| net.ReservoirSampling( |
| [self.reservoir, self.num_visited_blob, self.input_record.data(), |
| self.mutex] + self.extra_input_blobs, |
| [self.reservoir, self.num_visited_blob] + self.extra_output_blobs, |
| num_to_collect=self.num_to_collect, |
| ) |
|
|
