(self, model, input_record, num_to_collect,
name='reservoir_sampling', **kwargs)
| 18 | """ |
| 19 | |
| 20 | def __init__(self, model, input_record, num_to_collect, |
| 21 | name='reservoir_sampling', **kwargs): |
| 22 | super().__init__(model, name, input_record, **kwargs) |
| 23 | assert num_to_collect > 0 |
| 24 | self.num_to_collect = num_to_collect |
| 25 | |
| 26 | self.reservoir = self.create_param( |
| 27 | param_name='reservoir', |
| 28 | shape=[0], |
| 29 | initializer=('ConstantFill',), |
| 30 | optimizer=model.NoOptim, |
| 31 | ) |
| 32 | self.num_visited_blob = self.create_param( |
| 33 | param_name='num_visited', |
| 34 | shape=[], |
| 35 | initializer=('ConstantFill', { |
| 36 | 'value': 0, |
| 37 | 'dtype': core.DataType.INT64, |
| 38 | }), |
| 39 | optimizer=model.NoOptim, |
| 40 | ) |
| 41 | self.mutex = self.create_param( |
| 42 | param_name='mutex', |
| 43 | shape=[], |
| 44 | initializer=('CreateMutex',), |
| 45 | optimizer=model.NoOptim, |
| 46 | ) |
| 47 | |
| 48 | self.extra_input_blobs = [] |
| 49 | self.extra_output_blobs = [] |
| 50 | if 'object_id' in input_record: |
| 51 | object_to_pos = self.create_param( |
| 52 | param_name='object_to_pos', |
| 53 | shape=None, |
| 54 | initializer=('CreateMap', { |
| 55 | 'key_dtype': core.DataType.INT64, |
| 56 | 'valued_dtype': core.DataType.INT32, |
| 57 | }), |
| 58 | optimizer=model.NoOptim, |
| 59 | ) |
| 60 | pos_to_object = self.create_param( |
| 61 | param_name='pos_to_object', |
| 62 | shape=[0], |
| 63 | initializer=('ConstantFill', { |
| 64 | 'value': 0, |
| 65 | 'dtype': core.DataType.INT64, |
| 66 | }), |
| 67 | optimizer=model.NoOptim, |
| 68 | ) |
| 69 | self.extra_input_blobs.append(input_record.object_id()) |
| 70 | self.extra_input_blobs.extend([object_to_pos, pos_to_object]) |
| 71 | self.extra_output_blobs.extend([object_to_pos, pos_to_object]) |
| 72 | |
| 73 | self.output_schema = schema.Struct( |
| 74 | ( |
| 75 | 'reservoir', |
| 76 | schema.from_blob_list(input_record.data, [self.reservoir]) |
| 77 | ), |
nothing calls this directly
no test coverage detected