3 from __future__
import absolute_import
4 from __future__
import division
5 from __future__
import print_function
6 from __future__
import unicode_literals
10 from caffe2.proto
import caffe2_pb2
14 def lrn(model, blob_in, blob_out, order="NCHW", use_cudnn=False, **kwargs):
16 dev = kwargs[
'device_option']
if 'device_option' in kwargs \
17 else scope.CurrentDeviceScope()
18 is_cpu = dev
is None or dev.device_type == caffe2_pb2.CPU
19 if use_cudnn
and (
not is_cpu):
20 kwargs[
'engine'] =
'CUDNN' 23 blobs_out = [blob_out,
"_" + blob_out +
"_scale"]
31 if use_cudnn
and (
not is_cpu):
37 def softmax(model, blob_in, blob_out=None, use_cudnn=False, **kwargs):
40 kwargs[
'engine'] =
'CUDNN' 41 if blob_out
is not None:
42 return model.net.Softmax(blob_in, blob_out, **kwargs)
44 return model.net.Softmax(blob_in, **kwargs)
47 def instance_norm(model, blob_in, blob_out, dim_in, order="NCHW", **kwargs):
48 blob_out = blob_out
or model.net.NextName()
54 def init_blob(value, suffix):
55 return model.param_init_net.ConstantFill(
56 [], blob_out +
"_" + suffix, shape=[dim_in], value=value)
57 scale, bias = init_blob(1.0,
"s"), init_blob(0.0,
"b")
59 model.AddParameter(scale, ParameterTags.WEIGHT)
60 model.AddParameter(bias, ParameterTags.BIAS)
61 blob_outs = [blob_out, blob_out +
"_sm", blob_out +
"_siv"]
62 if 'is_test' in kwargs
and kwargs[
'is_test']:
63 blob_outputs = model.net.InstanceNorm(
64 [blob_in, scale, bias], [blob_out],
65 order=order, **kwargs)
68 blob_outputs = model.net.InstanceNorm(
69 [blob_in, scale, bias], blob_outs,
70 order=order, **kwargs)
72 return blob_outputs[0]
75 def spatial_bn(model, blob_in, blob_out, dim_in,
76 init_scale=1., init_bias=0.,
77 ScaleInitializer=
None, BiasInitializer=
None,
78 RunningMeanInitializer=
None, RunningVarianceInitializer=
None,
79 order=
"NCHW", **kwargs):
80 blob_out = blob_out
or model.net.NextName()
90 scale_init = (
"ConstantFill", {
'value': init_scale})
91 bias_init = (
"ConstantFill", {
'value': init_bias})
92 rm_init = (
"ConstantFill", {
'value': 0.0})
93 riv_init = (
"ConstantFill", {
'value': 1.0})
95 ScaleInitializer = initializers.update_initializer(
96 ScaleInitializer, scale_init, (
"ConstantFill", {})
98 BiasInitializer = initializers.update_initializer(
99 BiasInitializer, bias_init, (
"ConstantFill", {})
101 RunningMeanInitializer = initializers.update_initializer(
102 RunningMeanInitializer, rm_init, (
"ConstantFill", {})
104 RunningVarianceInitializer = initializers.update_initializer(
105 RunningVarianceInitializer, riv_init, (
"ConstantFill", {})
108 ScaleInitializer = initializers.ExternalInitializer()
109 BiasInitializer = initializers.ExternalInitializer()
110 RunningMeanInitializer = initializers.ExternalInitializer()
111 RunningVarianceInitializer = initializers.ExternalInitializer()
113 scale = model.create_param(
114 param_name=blob_out +
'_s',
116 initializer=ScaleInitializer,
117 tags=ParameterTags.WEIGHT
120 bias = model.create_param(
121 param_name=blob_out +
'_b',
123 initializer=BiasInitializer,
124 tags=ParameterTags.BIAS
127 running_mean = model.create_param(
128 param_name=blob_out +
'_rm',
130 initializer=RunningMeanInitializer,
131 tags=ParameterTags.COMPUTED_PARAM
134 running_inv_var = model.create_param(
135 param_name=blob_out +
'_riv',
137 initializer=RunningVarianceInitializer,
138 tags=ParameterTags.COMPUTED_PARAM
141 blob_outs = [blob_out, running_mean, running_inv_var,
142 blob_out +
"_sm", blob_out +
"_siv"]
143 if 'is_test' in kwargs
and kwargs[
'is_test']:
144 blob_outputs = model.net.SpatialBN(
145 [blob_in, scale, bias, blob_outs[1], blob_outs[2]], [blob_out],
146 order=order, **kwargs)
149 blob_outputs = model.net.SpatialBN(
150 [blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
151 order=order, **kwargs)
153 return blob_outputs[0]
167 Layer normalizes the input, cf. https://arxiv.org/pdf/1607.06450.pdf. 170 blob_in: The input blob to layer normalize. 171 blob_out: The layer normalized output blob. 172 dim_in: The dimension of the scale and bias. For example, if blob_in is 173 a 2D design matrix and axis is 1, this would be the number of 175 axis: (optional) The axis to normalize. Typically the feature axis. 177 epsilon: (optional) A small value used for numerical stability in 178 calculation. Defaults to 1e-4. 179 initial_scale: (optional) The initial value for the learned scale 180 parameter. Defaults to 1.0 181 initial_bias: (optional) The initial value for the learned bias 182 parameter of the layerwise standard deviation. Defaults to 0.0. 185 A 3-tuple consisting of: 186 - The layer normalized input blob. 187 - The mean of the input blob across the given axis. 188 - The standard deviation of the input blob acress the given axis. 194 normalized, mean, stdev = model.net.LayerNorm(
196 [blob_out, blob_out +
"_mean", blob_out +
"_stdev"],
202 scale = model.create_param(
203 param_name=
'{}_scale'.format(blob_out),
205 initializer=initializers.Initializer(
209 tags=ParameterTags.WEIGHT,
213 bias = model.create_param(
214 param_name=
'{}_bias'.format(blob_out),
216 initializer=initializers.Initializer(
220 tags=ParameterTags.BIAS,
223 scaled = model.net.Mul(
225 [
'{}_scaled'.format(blob_out)],
230 biased = model.net.Add(
232 [
'{}_biased'.format(blob_out)],
237 return biased, mean, stdev