Caffe2 - Python API
A deep learning, cross platform ML framework
batch_normalization.py
1 # Copyright (c) 2016-present, Facebook, Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 ##############################################################################
15 
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 from __future__ import unicode_literals
20 
21 from caffe2.python import schema
22 from caffe2.python.layers.layers import ModelLayer
23 
24 import numpy as np
25 
26 
28  def __init__(
29  self,
30  model,
31  input_record,
32  name='batch_normalization',
33  scale_optim=None,
34  bias_optim=None,
35  momentum=0.9,
36  order='NCHW',
37  **kwargs
38  ):
39  super(BatchNormalization, self).__init__(
40  model, name, input_record, **kwargs)
41 
42  assert isinstance(input_record, schema.Scalar), "Incorrect input type"
43 
44  self.input_shape = input_record.field_type().shape
45 
46  if len(self.input_shape) == 3:
47  if order == "NCHW":
48  input_dims = self.input_shape[0]
49  elif order == "NHWC":
50  input_dims = self.input_shape[2]
51  else:
52  raise ValueError("Please specify a correct order")
53  else:
54  assert len(self.input_shape) == 1, (
55  "This layer supports only 4D or 2D tesnors")
56  input_dims = self.input_shape[0]
57 
59  (np.float32, self.input_shape),
60  self.get_next_blob_reference('output')
61  )
62 
63  self.momentum = momentum
64  self.order = order
65 
66  self.scale = self.create_param(param_name='scale',
67  shape=[input_dims],
68  initializer=('ConstantFill', {'value': 1.0}),
69  optimizer=scale_optim)
70  self.bias = self.create_param(param_name='bias',
71  shape=[input_dims],
72  initializer=('ConstantFill', {'value': 0.0}),
73  optimizer=bias_optim)
74  self.rm = self.create_param(param_name='running_mean',
75  shape=[input_dims],
76  initializer=('ConstantFill', {'value': 0.0}),
77  optimizer=model.NoOptim)
78  self.riv = self.create_param(param_name='running_inv_var',
79  shape=[input_dims],
80  initializer=('ConstantFill', {'value': 1.0}),
81  optimizer=model.NoOptim)
82 
83  def _add_ops(self, net, is_test, out_blob=None):
84  original_input_blob = self.input_record.field_blobs()
85  input_blob = net.NextScopedBlob('expand_input')
86  if len(self.input_shape) == 1:
87  input_blob = net.ExpandDims(original_input_blob,
88  dims=[2, 3])
89  else:
90  input_blob = original_input_blob[0]
91 
92  if out_blob is None:
93  bn_output = self.output_schema.field_blobs()
94  else:
95  bn_output = out_blob
96  if is_test:
97  output_blobs = bn_output
98  else:
99  output_blobs = bn_output + [self.rm, self.riv,
100  net.NextScopedBlob('bn_saved_mean'),
101  net.NextScopedBlob('bn_saved_iv')]
102 
103  net.SpatialBN([input_blob, self.scale,
104  self.bias, self.rm, self.riv],
105  output_blobs,
106  momentum=self.momentum,
107  is_test=is_test,
108  order=self.order)
109 
110  if len(self.input_shape) == 1:
111  net.Squeeze(bn_output,
112  bn_output,
113  dims=[2, 3])
114 
115  def add_train_ops(self, net):
116  self._add_ops(net, is_test=False)
117 
118  def add_eval_ops(self, net):
119  self._add_ops(net, is_test=True)
120 
121  def add_ops(self, net):
122  self.add_eval_ops(net)
def get_next_blob_reference(self, name)
Definition: layers.py:352
def create_param(self, param_name, shape, initializer, optimizer, ps_param=None, regularizer=None)
Definition: layers.py:337