|
|
|
|
|
|
|
|
|
|
| from caffe2.python import schema |
| from caffe2.python.layers.layers import ModelLayer |
|
|
| import numpy as np |
|
|
|
|
| class LayerNormalization(ModelLayer): |
| def __init__( |
| self, |
| model, |
| input_record, |
| name='layer_normalization', |
| scale_optim=None, |
| bias_optim=None, |
| epsilon=1e-4, |
| axis=1, |
| use_layer_norm_op=True, |
| scale_init_value=1.0, |
| **kwargs |
| ): |
| super(LayerNormalization, self).__init__( |
| model, name, input_record, **kwargs) |
|
|
| assert isinstance(input_record, schema.Scalar), ( |
| "Incorrect input type: {}".format(input_record)) |
|
|
| self.input_shape = input_record.field_type().shape |
| self.axis = axis |
|
|
| assert len(self.input_shape) >= 1, ( |
| "This layer supports only >= 2D tensors") |
| input_dims = self.input_shape[0] |
|
|
| self.output_schema = schema.Scalar( |
| (np.float32, self.input_shape), |
| self.get_next_blob_reference('output') |
| ) |
|
|
| self.scale = self.create_param(param_name='scale', |
| shape=[input_dims], |
| initializer=('ConstantFill', {'value': scale_init_value}), |
| optimizer=scale_optim) |
| self.bias = self.create_param(param_name='bias', |
| shape=[input_dims], |
| initializer=('ConstantFill', {'value': 0.0}), |
| optimizer=bias_optim) |
| self.use_layer_norm_op = use_layer_norm_op |
|
|
| if self.use_layer_norm_op: |
| self.epsilon = epsilon |
| else: |
| assert len(self.input_shape) == 1, ( |
| "When using alternative implementation, " |
| "input data can only be 2D" |
| ) |
| self.epsilon = model.maybe_add_global_constant( |
| "%s_epsilon" % self.name, float(epsilon) |
| ) |
|
|
| def add_ops_with_layer_norm_op(self, net): |
| input_blob = self.input_record.field_blobs() |
| ln_output = self.output_schema.field_blobs() |
|
|
| output_blobs = [net.NextScopedBlob('ln_output'), net.NextScopedBlob('ln_mean'), |
| net.NextScopedBlob('ln_stdev')] |
|
|
| normalized, mean, stdev = net.LayerNorm(input_blob, |
| output_blobs, |
| axis=self.axis, |
| epsilon=self.epsilon) |
|
|
| scaled = net.Mul( |
| [normalized, self.scale], |
| [net.NextScopedBlob('ln_scaled')], |
| broadcast=1, |
| axis=self.axis, |
| ) |
|
|
| net.Add( |
| [scaled, self.bias], |
| ln_output, |
| broadcast=1, |
| axis=self.axis, |
| ) |
|
|
| def add_ops_without_layer_norm_op(self, net): |
| |
| |
| |
| ln_output = net.NextScopedBlob("ln_output") |
| ln_mean = net.NextScopedBlob("ln_mean") |
| ln_stdev = net.NextScopedBlob("ln_stdev") |
| ln_mean_arr = net.NextScopedBlob("ln_mean_arr") |
| net.ReduceBackMean(self.input_record.field_blobs(), [ln_mean_arr]) |
| net.ExpandDims([ln_mean_arr], [ln_mean], dims=[1]) |
| ln_centered = net.NextScopedBlob("ln_centered") |
| net.Sub(self.input_record.field_blobs() + [ln_mean], [ln_centered]) |
| ln_sqr = net.NextScopedBlob("ln_sqr") |
| net.Sqr([ln_centered], [ln_sqr]) |
| ln_sqr_mean = net.NextScopedBlob("ln_sqr_mean") |
| net.ReduceBackMean([ln_sqr], [ln_sqr_mean]) |
| ln_var = net.NextScopedBlob("ln_var") |
| net.Add([ln_sqr_mean, self.epsilon], ln_var) |
| ln_std_arr = net.NextScopedBlob("ln_std_arr") |
| net.Pow([ln_var], [ln_std_arr], exponent=0.5) |
| net.ExpandDims([ln_std_arr], [ln_stdev], dims=[1]) |
| net.Div([ln_centered, ln_stdev], [ln_output]) |
| ln_scaled = net.NextScopedBlob("ln_scaled") |
| net.Mul([ln_output, self.scale], [ln_scaled]) |
| net.Add([ln_scaled, self.bias], self.output_schema.field_blobs()) |
|
|
| def add_ops(self, net): |
| if self.use_layer_norm_op: |
| self.add_ops_with_layer_norm_op(net) |
| else: |
| self.add_ops_without_layer_norm_op(net) |
|
|
