| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ModelLayer |
| | import math |
| |
|
| |
|
| | class AddBias(ModelLayer): |
| |
|
| | def __init__(self, model, input_record, bias_init=None, |
| | bias_optim=None, name='add_bias'): |
| | super(AddBias, self).__init__(model, name, input_record) |
| | assert isinstance(input_record, schema.Scalar), "Incorrect input type" |
| | assert len(input_record.field_type().shape) > 0, ( |
| | "AddBias expects limited dimensions of the input tensor") |
| |
|
| | input_dims = input_record.field_type().shape[0] |
| | assert input_dims > 0, ( |
| | "AddBias expects input dimensions > 0, got {}".format(input_dims)) |
| |
|
| | scale = math.sqrt(1.0 / input_dims) |
| | bias_init = bias_init if bias_init else ( |
| | 'UniformFill', {'min': -scale, 'max': scale}) |
| |
|
| | self.b = self.create_param( |
| | param_name='b', |
| | shape=[input_dims, ], |
| | initializer=bias_init, |
| | optimizer=bias_optim, |
| | ) |
| |
|
| | self.output_schema = schema.Scalar( |
| | (input_record.field_type().base, (input_dims, )), |
| | self.get_next_blob_reference('output') |
| | ) |
| |
|
| | def add_ops(self, net): |
| | net.Add(self.input_record.field_blobs() + [self.b], |
| | self.output_schema.field_blobs(), broadcast=1) |
| |
|