| |
| |
|
|
|
|
|
|
|
|
|
|
| from caffe2.python import schema |
| from caffe2.python.layers.layers import ModelLayer |
| import numpy as np |
|
|
|
|
| class ConstantWeight(ModelLayer): |
| def __init__( |
| self, |
| model, |
| input_record, |
| weights=None, |
| name='constant_weight', |
| **kwargs |
| ): |
| super(ConstantWeight, |
| self).__init__(model, name, input_record, **kwargs) |
| self.output_schema = schema.Scalar( |
| np.float32, self.get_next_blob_reference('constant_weight') |
| ) |
| self.data = self.input_record.field_blobs() |
| self.num = len(self.data) |
| weights = ( |
| weights if weights is not None else |
| [1. / self.num for _ in range(self.num)] |
| ) |
| assert len(weights) == self.num |
| self.weights = [ |
| self.model.add_global_constant( |
| '%s_weight_%d' % (self.name, i), float(weights[i]) |
| ) for i in range(self.num) |
| ] |
|
|
| def add_ops(self, net): |
| net.WeightedSum( |
| [b for x_w_pair in zip(self.data, self.weights) for b in x_w_pair], |
| self.output_schema() |
| ) |
|
|
