Caffe2 - Python API
A deep learning, cross platform ML framework
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 ## @package normalization
17 # Module caffe2.python.helpers.normalization
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 from __future__ import unicode_literals
22 
23 from caffe2.python import scope
24 from caffe2.python.modeling.parameter_info import ParameterTags
25 from caffe2.proto import caffe2_pb2
26 from caffe2.python.modeling import initializers
27 
28 
29 def lrn(model, blob_in, blob_out, order="NCHW", use_cudnn=False, **kwargs):
30  """LRN"""
31  dev = kwargs['device_option'] if 'device_option' in kwargs \
32  else scope.CurrentDeviceScope()
33  is_cpu = dev is None or dev.device_type == caffe2_pb2.CPU
34  if use_cudnn and (not is_cpu):
35  kwargs['engine'] = 'CUDNN'
36  blobs_out = blob_out
37  else:
38  blobs_out = [blob_out, "_" + blob_out + "_scale"]
39  lrn = model.net.LRN(
40  blob_in,
41  blobs_out,
42  order=order,
43  **kwargs
44  )
45 
46  if use_cudnn and (not is_cpu):
47  return lrn
48  else:
49  return lrn[0]
50 
51 
52 def softmax(model, blob_in, blob_out=None, use_cudnn=False, **kwargs):
53  """Softmax."""
54  if use_cudnn:
55  kwargs['engine'] = 'CUDNN'
56  if blob_out is not None:
57  return model.net.Softmax(blob_in, blob_out, **kwargs)
58  else:
59  return model.net.Softmax(blob_in, **kwargs)
60 
61 
62 def instance_norm(model, blob_in, blob_out, dim_in, order="NCHW", **kwargs):
63  blob_out = blob_out or model.net.NextName()
64  # Input: input, scale, bias
65  # Output: output, saved_mean, saved_inv_std
66  # scale: initialize with ones
67  # bias: initialize with zeros
68 
69  def init_blob(value, suffix):
70  return model.param_init_net.ConstantFill(
71  [], blob_out + "_" + suffix, shape=[dim_in], value=value)
72  scale, bias = init_blob(1.0, "s"), init_blob(0.0, "b")
73 
74  model.AddParameter(scale, ParameterTags.WEIGHT)
75  model.AddParameter(bias, ParameterTags.BIAS)
76  blob_outs = [blob_out, blob_out + "_sm", blob_out + "_siv"]
77  if 'is_test' in kwargs and kwargs['is_test']:
78  blob_outputs = model.net.InstanceNorm(
79  [blob_in, scale, bias], [blob_out],
80  order=order, **kwargs)
81  return blob_outputs
82  else:
83  blob_outputs = model.net.InstanceNorm(
84  [blob_in, scale, bias], blob_outs,
85  order=order, **kwargs)
86  # Return the output
87  return blob_outputs[0]
88 
89 
90 def spatial_bn(model, blob_in, blob_out, dim_in,
91  init_scale=1., init_bias=0.,
92  ScaleInitializer=None, BiasInitializer=None,
93  RunningMeanInitializer=None, RunningVarianceInitializer=None,
94  order="NCHW", **kwargs):
95  blob_out = blob_out or model.net.NextName()
96  # Input: input, scale, bias, est_mean, est_inv_var
97  # Output: output, running_mean, running_inv_var, saved_mean,
98  # saved_inv_var
99  # scale: initialize with init_scale (default 1.)
100  # bias: initialize with init_bias (default 0.)
101  # est mean: zero
102  # est var: ones
103 
104  if model.init_params:
105  scale_init = ("ConstantFill", {'value': init_scale})
106  bias_init = ("ConstantFill", {'value': init_bias})
107  rm_init = ("ConstantFill", {'value': 0.0})
108  riv_init = ("ConstantFill", {'value': 1.0})
109 
110  ScaleInitializer = initializers.update_initializer(
111  ScaleInitializer, scale_init, ("ConstantFill", {})
112  )
113  BiasInitializer = initializers.update_initializer(
114  BiasInitializer, bias_init, ("ConstantFill", {})
115  )
116  RunningMeanInitializer = initializers.update_initializer(
117  RunningMeanInitializer, rm_init, ("ConstantFill", {})
118  )
119  RunningVarianceInitializer = initializers.update_initializer(
120  RunningVarianceInitializer, riv_init, ("ConstantFill", {})
121  )
122  else:
123  ScaleInitializer = initializers.ExternalInitializer()
124  BiasInitializer = initializers.ExternalInitializer()
125  RunningMeanInitializer = initializers.ExternalInitializer()
126  RunningVarianceInitializer = initializers.ExternalInitializer()
127 
128  scale = model.create_param(
129  param_name=blob_out + '_s',
130  shape=[dim_in],
131  initializer=ScaleInitializer,
132  tags=ParameterTags.WEIGHT
133  )
134 
135  bias = model.create_param(
136  param_name=blob_out + '_b',
137  shape=[dim_in],
138  initializer=BiasInitializer,
139  tags=ParameterTags.BIAS
140  )
141 
142  running_mean = model.create_param(
143  param_name=blob_out + '_rm',
144  shape=[dim_in],
145  initializer=RunningMeanInitializer,
146  tags=ParameterTags.COMPUTED_PARAM
147  )
148 
149  running_inv_var = model.create_param(
150  param_name=blob_out + '_riv',
151  shape=[dim_in],
152  initializer=RunningVarianceInitializer,
153  tags=ParameterTags.COMPUTED_PARAM
154  )
155 
156  blob_outs = [blob_out, running_mean, running_inv_var,
157  blob_out + "_sm", blob_out + "_siv"]
158  if 'is_test' in kwargs and kwargs['is_test']:
159  blob_outputs = model.net.SpatialBN(
160  [blob_in, scale, bias, blob_outs[1], blob_outs[2]], [blob_out],
161  order=order, **kwargs)
162  return blob_outputs
163  else:
164  blob_outputs = model.net.SpatialBN(
165  [blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
166  order=order, **kwargs)
167  # Return the output
168  return blob_outputs[0]
169 
170 
171 def layer_norm(
172  model,
173  blob_in,
174  blob_out,
175  dim_in,
176  axis=1,
177  epsilon=1e-4,
178  initial_scale=1.0,
179  initial_bias=0.0,
180 ):
181  '''
182  Layer normalizes the input, cf. https://arxiv.org/pdf/1607.06450.pdf.
183 
184  Args:
185  blob_in: The input blob to layer normalize.
186  blob_out: The layer normalized output blob.
187  dim_in: The dimension of the scale and bias. For example, if blob_in is
188  a 2D design matrix and axis is 1, this would be the number of
189  columns.
190  axis: (optional) The axis to normalize. Typically the feature axis.
191  Defaults to 1.
192  epsilon: (optional) A small value used for numerical stability in
193  calculation. Defaults to 1e-4.
194  initial_scale: (optional) The initial value for the learned scale
195  parameter. Defaults to 1.0
196  initial_bias: (optional) The initial value for the learned bias
197  parameter of the layerwise standard deviation. Defaults to 0.0.
198 
199  Returns:
200  A 3-tuple consisting of:
201  - The layer normalized input blob.
202  - The mean of the input blob across the given axis.
203  - The standard deviation of the input blob acress the given axis.
204  '''
205 
206  # The LayerNorm operator only performs the layerwise z-shift, without
207  # scaling and shifting by the learned scale and bias parameters. We have
208  # to do that separately below.
209  normalized, mean, stdev = model.net.LayerNorm(
210  [blob_in],
211  [blob_out, blob_out + "_mean", blob_out + "_stdev"],
212  axis=axis,
213  epsilon=epsilon,
214  )
215 
216  # The learned multiplicative scale or "gain".
217  scale = model.create_param(
218  param_name='{}_scale'.format(blob_out),
219  shape=[dim_in],
220  initializer=initializers.Initializer(
221  'ConstantFill',
222  value=initial_scale,
223  ),
224  tags=ParameterTags.WEIGHT,
225  )
226 
227  # The learned additive bias or "shift".
228  bias = model.create_param(
229  param_name='{}_bias'.format(blob_out),
230  shape=[dim_in],
231  initializer=initializers.Initializer(
232  'ConstantFill',
233  value=initial_bias,
234  ),
235  tags=ParameterTags.BIAS,
236  )
237 
238  scaled = model.net.Mul(
239  [normalized, scale],
240  ['{}_scaled'.format(blob_out)],
241  broadcast=1,
242  axis=axis,
243  )
244 
245  biased = model.net.Add(
246  [scaled, bias],
247  ['{}_biased'.format(blob_out)],
248  broadcast=1,
249  axis=axis,
250  )
251 
252  return biased, mean, stdev