| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ModelLayer |
| |
|
| |
|
| | class BlobWeightedSum(ModelLayer): |
| | """ |
| | This layer implements the weighted sum: |
| | weighted element-wise sum of input blobs. |
| | """ |
| | def __init__( |
| | self, |
| | model, |
| | input_record, |
| | init_weights=None, |
| | weight_optim=None, |
| | name='blob_weighted_sum', |
| | **kwargs |
| | ): |
| | super(BlobWeightedSum, self).__init__(model, name, input_record, **kwargs) |
| |
|
| | self.blobs = self.input_record.field_blobs() |
| |
|
| | self.num_weights = len(self.blobs) |
| | assert self.num_weights > 1, ( |
| | "BlobWeightedSum expects more than one input blobs" |
| | ) |
| |
|
| | assert len(input_record.field_types()[0].shape) > 0, ( |
| | "BlobWeightedSum expects limited dimensions of the input tensor" |
| | ) |
| |
|
| | assert all( |
| | input_record.field_types()[0].shape == input_record.field_types()[i].shape |
| | for i in range(1, self.num_weights) |
| | ), "Shape of input blobs should be the same shape {}".format( |
| | input_record.field_types()[0].shape |
| | ) |
| |
|
| | if init_weights: |
| | assert self.num_weights == len(init_weights), ( |
| | "the size of init_weights should be the same as input blobs, " |
| | "expects {}, got {}".format(self.num_weights, len(init_weights)) |
| | ) |
| | else: |
| | init_weights = [1.0] * self.num_weights |
| |
|
| | self.weights = [ |
| | self.create_param( |
| | param_name="w_{}".format(idx), |
| | shape=[1], |
| | initializer=('ConstantFill', {'value': float(init_weights[idx])}), |
| | optimizer=weight_optim |
| | ) for idx in range(self.num_weights) |
| | ] |
| |
|
| | self.output_schema = schema.Scalar( |
| | input_record.field_types()[0], |
| | self.get_next_blob_reference('blob_weighted_sum_out') |
| | ) |
| |
|
| | def add_ops(self, net): |
| | net.WeightedSum( |
| | [x for pair in zip(self.blobs, self.weights) for x in pair], |
| | self.output_schema(), |
| | grad_on_w=True, |
| | ) |
| |
|
