caffe2/python/helpers/normalization.py - platform/external/pytorch - Git at Google

 ## @package normalization
 # Module caffe2.python.helpers.normalization


 from caffe2.python import scope
 from caffe2.python.modeling.parameter_info import ParameterTags
 from caffe2.proto import caffe2_pb2
 from caffe2.python.modeling import initializers


 def lrn(model, blob_in, blob_out, order="NCHW", use_cudnn=False, **kwargs):
     """LRN"""
     dev = kwargs['device_option'] if 'device_option' in kwargs \
         else scope.CurrentDeviceScope()
     is_cpu = dev is None or dev.device_type == caffe2_pb2.CPU
     if use_cudnn and (not is_cpu):
         kwargs['engine'] = 'CUDNN'
         blobs_out = blob_out
     else:
         blobs_out = [blob_out, "_" + blob_out + "_scale"]
     lrn = model.net.LRN(
         blob_in,
         blobs_out,
         order=order,
         **kwargs
     )

     if use_cudnn and (not is_cpu):
         return lrn
     else:
         return lrn[0]


 def softmax(model, blob_in, blob_out=None, use_cudnn=False, **kwargs):
     """Softmax."""
     if use_cudnn:
         kwargs['engine'] = 'CUDNN'
     if blob_out is not None:
         return model.net.Softmax(blob_in, blob_out, **kwargs)
     else:
         return model.net.Softmax(blob_in, **kwargs)


 def instance_norm(model, blob_in, blob_out, dim_in, order="NCHW", **kwargs):
     blob_out = blob_out or model.net.NextName()
     # Input: input, scale, bias
     # Output: output, saved_mean, saved_inv_std
     # scale: initialize with ones
     # bias: initialize with zeros

     def init_blob(value, suffix):
         return model.param_init_net.ConstantFill(
             [], blob_out + "_" + suffix, shape=[dim_in], value=value)
     scale, bias = init_blob(1.0, "s"), init_blob(0.0, "b")

     model.AddParameter(scale, ParameterTags.WEIGHT)
     model.AddParameter(bias, ParameterTags.BIAS)
     blob_outs = [blob_out, blob_out + "_sm", blob_out + "_siv"]
     if 'is_test' in kwargs and kwargs['is_test']:
         blob_outputs = model.net.InstanceNorm(
             [blob_in, scale, bias], [blob_out],
             order=order, **kwargs)
         return blob_outputs
     else:
         blob_outputs = model.net.InstanceNorm(
             [blob_in, scale, bias], blob_outs,
             order=order, **kwargs)
         # Return the output
         return blob_outputs[0]


 def spatial_bn(model, blob_in, blob_out, dim_in,
                init_scale=1., init_bias=0.,
                ScaleInitializer=None, BiasInitializer=None,
                RunningMeanInitializer=None, RunningVarianceInitializer=None,
                order="NCHW", **kwargs):
     blob_out = blob_out or model.net.NextName()
     # Input: input, scale, bias, est_mean, est_inv_var
     # Output: output, running_mean, running_inv_var, saved_mean,
     #         saved_inv_var
     # scale: initialize with init_scale (default 1.)
     # bias: initialize with init_bias (default 0.)
     # est mean: zero
     # est var: ones

     if model.init_params:
         scale_init = ("ConstantFill", {'value': init_scale})
         bias_init = ("ConstantFill", {'value': init_bias})
         rm_init = ("ConstantFill", {'value': 0.0})
         riv_init = ("ConstantFill", {'value': 1.0})

         ScaleInitializer = initializers.update_initializer(
             ScaleInitializer, scale_init, ("ConstantFill", {})
         )
         BiasInitializer = initializers.update_initializer(
             BiasInitializer, bias_init, ("ConstantFill", {})
         )
         RunningMeanInitializer = initializers.update_initializer(
             RunningMeanInitializer, rm_init, ("ConstantFill", {})
         )
         RunningVarianceInitializer = initializers.update_initializer(
             RunningVarianceInitializer, riv_init, ("ConstantFill", {})
         )
     else:
         ScaleInitializer = initializers.ExternalInitializer()
         BiasInitializer = initializers.ExternalInitializer()
         RunningMeanInitializer = initializers.ExternalInitializer()
         RunningVarianceInitializer = initializers.ExternalInitializer()

     scale = model.create_param(
         param_name=blob_out + '_s',
         shape=[dim_in],
         initializer=ScaleInitializer,
         tags=ParameterTags.WEIGHT
     )

     bias = model.create_param(
         param_name=blob_out + '_b',
         shape=[dim_in],
         initializer=BiasInitializer,
         tags=ParameterTags.BIAS
     )

     running_mean = model.create_param(
         param_name=blob_out + '_rm',
         shape=[dim_in],
         initializer=RunningMeanInitializer,
         tags=ParameterTags.COMPUTED_PARAM
     )

     running_inv_var = model.create_param(
         param_name=blob_out + '_riv',
         shape=[dim_in],
         initializer=RunningVarianceInitializer,
         tags=ParameterTags.COMPUTED_PARAM
     )

     blob_outs = [blob_out, running_mean, running_inv_var,
                  blob_out + "_sm", blob_out + "_siv"]
     if 'is_test' in kwargs and kwargs['is_test']:
         blob_outputs = model.net.SpatialBN(
             [blob_in, scale, bias, blob_outs[1], blob_outs[2]], [blob_out],
             order=order, **kwargs)
         return blob_outputs
     else:
         blob_outputs = model.net.SpatialBN(
             [blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
             order=order, **kwargs)
         # Return the output
         return blob_outputs[0]


 def spatial_gn(model, blob_in, blob_out, dim_in,
                init_scale=1., init_bias=0.,
                ScaleInitializer=None, BiasInitializer=None,
                RunningMeanInitializer=None, RunningVarianceInitializer=None,
                order="NCHW", **kwargs):
     '''
     Group normalizes the input, cf. https://arxiv.org/abs/1803.08494.
     '''

     blob_out = blob_out or model.net.NextName()
     # Input: input, scale, bias
     # Output: output, group_mean, group_inv_std
     # scale: initialize with init_scale (default 1.)
     # [recommendation: set init_scale = 0. in the last layer for each res block]
     # bias: initialize with init_bias (default 0.)

     if model.init_params:
         scale_init = ("ConstantFill", {'value': init_scale})
         bias_init = ("ConstantFill", {'value': init_bias})

         ScaleInitializer = initializers.update_initializer(
             ScaleInitializer, scale_init, ("ConstantFill", {})
         )
         BiasInitializer = initializers.update_initializer(
             BiasInitializer, bias_init, ("ConstantFill", {})
         )
     else:
         ScaleInitializer = initializers.ExternalInitializer()
         BiasInitializer = initializers.ExternalInitializer()

     scale = model.create_param(
         param_name=blob_out + '_s',
         shape=[dim_in],
         initializer=ScaleInitializer,
         tags=ParameterTags.WEIGHT
     )

     bias = model.create_param(
         param_name=blob_out + '_b',
         shape=[dim_in],
         initializer=BiasInitializer,
         tags=ParameterTags.BIAS
     )

     blob_outs = [blob_out,
                  blob_out + "_mean", blob_out + "_std"]

     blob_outputs = model.net.GroupNorm(
         [blob_in, scale, bias],
         blob_outs,
         **kwargs)
     # Return the output
     return blob_outputs[0]


 def layer_norm(
     model,
     blob_in,
     blob_out,
     dim_in,
     axis=1,
     epsilon=1e-4,
     initial_scale=1.0,
     initial_bias=0.0,
 ):
     '''
     Layer normalizes the input, cf. https://arxiv.org/pdf/1607.06450.pdf.

     Args:
         blob_in: The input blob to layer normalize.
         blob_out: The layer normalized output blob.
         dim_in: The dimension of the scale and bias. For example, if blob_in is
             a 2D design matrix and axis is 1, this would be the number of
             columns.
         axis: (optional) The axis to normalize. Typically the feature axis.
             Defaults to 1.
         epsilon: (optional) A small value used for numerical stability in
             calculation. Defaults to 1e-4.
         initial_scale: (optional) The initial value for the learned scale
             parameter. Defaults to 1.0
         initial_bias: (optional) The initial value for the learned bias
             parameter of the layerwise standard deviation. Defaults to 0.0.

     Returns:
         A 3-tuple consisting of:
             - The layer normalized input blob.
             - The mean of the input blob across the given axis.
             - The standard deviation of the input blob acress the given axis.
     '''

     # The learned multiplicative scale or "gain".
     scale = model.create_param(
         param_name='{}_scale'.format(blob_out),
         shape=[dim_in] if isinstance(dim_in, int) else dim_in,
         initializer=initializers.Initializer(
             'ConstantFill',
             value=initial_scale,
         ),
         tags=ParameterTags.WEIGHT,
     )

     # The learned additive bias or "shift".
     bias = model.create_param(
         param_name='{}_bias'.format(blob_out),
         shape=[dim_in] if isinstance(dim_in, int) else dim_in,
         initializer=initializers.Initializer(
             'ConstantFill',
             value=initial_bias,
         ),
         tags=ParameterTags.BIAS,
     )

     normalized, mean, std = model.net.LayerNorm(
         [blob_in, scale, bias],
         [blob_out, blob_out + "_mean", blob_out + "_std"],
         axis=axis,
         epsilon=epsilon,
         elementwise_affine=True,
     )

     return normalized, mean, std

 def moments_with_running_stats(model, blob_in, blob_out, dim_in,
                                      RunningMeanInitializer=None, RunningVarianceInitializer=None,
                                      order="NCHW", **kwargs):

     if model.init_params:
         rm_init = ("ConstantFill", {'value': 0.0})
         riv_init = ("ConstantFill", {'value': 1.0})

         RunningMeanInitializer = initializers.update_initializer(
             RunningMeanInitializer, rm_init, ("ConstantFill", {})
         )
         RunningVarianceInitializer = initializers.update_initializer(
             RunningVarianceInitializer, riv_init, ("ConstantFill", {})
         )
     else:
         RunningMeanInitializer = initializers.ExternalInitializer()
         RunningVarianceInitializer = initializers.ExternalInitializer()

     running_mean = model.create_param(
         param_name=blob_out + '_rm',
         shape=[dim_in],
         initializer=RunningMeanInitializer,
         tags=ParameterTags.COMPUTED_PARAM
     )

     # this is just running variance
     running_inv_var = model.create_param(
         param_name=blob_out + '_riv',
         shape=[dim_in],
         initializer=RunningVarianceInitializer,
         tags=ParameterTags.COMPUTED_PARAM
     )

     blob_outs = [blob_out + "_sm", blob_out + "_sv"]
     if order == 'NCHW':
         blob_outputs = model.net.Moments(
             [blob_in], blob_outs,
             axes=[0, 2, 3],
             order=order, keepdims=False, **kwargs)
     elif order == 'NHWC':
         blob_outputs = model.net.Moments(
             [blob_in], blob_outs,
             axes=[0, 1, 2],
             order=order, keepdims=False, **kwargs)
     return blob_outputs
	## @package normalization
	# Module caffe2.python.helpers.normalization





	from caffe2.python import scope
	from caffe2.python.modeling.parameter_info import ParameterTags
	from caffe2.proto import caffe2_pb2
	from caffe2.python.modeling import initializers


	def lrn(model, blob_in, blob_out, order="NCHW", use_cudnn=False, **kwargs):
	"""LRN"""
	dev = kwargs['device_option'] if 'device_option' in kwargs \
	else scope.CurrentDeviceScope()
	is_cpu = dev is None or dev.device_type == caffe2_pb2.CPU
	if use_cudnn and (not is_cpu):
	kwargs['engine'] = 'CUDNN'
	blobs_out = blob_out
	else:
	blobs_out = [blob_out, "_" + blob_out + "_scale"]
	lrn = model.net.LRN(
	blob_in,
	blobs_out,
	order=order,
	**kwargs
	)

	if use_cudnn and (not is_cpu):
	return lrn
	else:
	return lrn[0]


	def softmax(model, blob_in, blob_out=None, use_cudnn=False, **kwargs):
	"""Softmax."""
	if use_cudnn:
	kwargs['engine'] = 'CUDNN'
	if blob_out is not None:
	return model.net.Softmax(blob_in, blob_out, **kwargs)
	else:
	return model.net.Softmax(blob_in, **kwargs)


	def instance_norm(model, blob_in, blob_out, dim_in, order="NCHW", **kwargs):
	blob_out = blob_out or model.net.NextName()
	# Input: input, scale, bias
	# Output: output, saved_mean, saved_inv_std
	# scale: initialize with ones
	# bias: initialize with zeros

	def init_blob(value, suffix):
	return model.param_init_net.ConstantFill(
	[], blob_out + "_" + suffix, shape=[dim_in], value=value)
	scale, bias = init_blob(1.0, "s"), init_blob(0.0, "b")

	model.AddParameter(scale, ParameterTags.WEIGHT)
	model.AddParameter(bias, ParameterTags.BIAS)
	blob_outs = [blob_out, blob_out + "_sm", blob_out + "_siv"]
	if 'is_test' in kwargs and kwargs['is_test']:
	blob_outputs = model.net.InstanceNorm(
	[blob_in, scale, bias], [blob_out],
	order=order, **kwargs)
	return blob_outputs
	else:
	blob_outputs = model.net.InstanceNorm(
	[blob_in, scale, bias], blob_outs,
	order=order, **kwargs)
	# Return the output
	return blob_outputs[0]


	def spatial_bn(model, blob_in, blob_out, dim_in,
	init_scale=1., init_bias=0.,
	ScaleInitializer=None, BiasInitializer=None,
	RunningMeanInitializer=None, RunningVarianceInitializer=None,
	order="NCHW", **kwargs):
	blob_out = blob_out or model.net.NextName()
	# Input: input, scale, bias, est_mean, est_inv_var
	# Output: output, running_mean, running_inv_var, saved_mean,
	# saved_inv_var
	# scale: initialize with init_scale (default 1.)
	# bias: initialize with init_bias (default 0.)
	# est mean: zero
	# est var: ones

	if model.init_params:
	scale_init = ("ConstantFill", {'value': init_scale})
	bias_init = ("ConstantFill", {'value': init_bias})
	rm_init = ("ConstantFill", {'value': 0.0})
	riv_init = ("ConstantFill", {'value': 1.0})

	ScaleInitializer = initializers.update_initializer(
	ScaleInitializer, scale_init, ("ConstantFill", {})
	)
	BiasInitializer = initializers.update_initializer(
	BiasInitializer, bias_init, ("ConstantFill", {})
	)
	RunningMeanInitializer = initializers.update_initializer(
	RunningMeanInitializer, rm_init, ("ConstantFill", {})
	)
	RunningVarianceInitializer = initializers.update_initializer(
	RunningVarianceInitializer, riv_init, ("ConstantFill", {})
	)
	else:
	ScaleInitializer = initializers.ExternalInitializer()
	BiasInitializer = initializers.ExternalInitializer()
	RunningMeanInitializer = initializers.ExternalInitializer()
	RunningVarianceInitializer = initializers.ExternalInitializer()

	scale = model.create_param(
	param_name=blob_out + '_s',
	shape=[dim_in],
	initializer=ScaleInitializer,
	tags=ParameterTags.WEIGHT
	)

	bias = model.create_param(
	param_name=blob_out + '_b',
	shape=[dim_in],
	initializer=BiasInitializer,
	tags=ParameterTags.BIAS
	)

	running_mean = model.create_param(
	param_name=blob_out + '_rm',
	shape=[dim_in],
	initializer=RunningMeanInitializer,
	tags=ParameterTags.COMPUTED_PARAM
	)

	running_inv_var = model.create_param(
	param_name=blob_out + '_riv',
	shape=[dim_in],
	initializer=RunningVarianceInitializer,
	tags=ParameterTags.COMPUTED_PARAM
	)

	blob_outs = [blob_out, running_mean, running_inv_var,
	blob_out + "_sm", blob_out + "_siv"]
	if 'is_test' in kwargs and kwargs['is_test']:
	blob_outputs = model.net.SpatialBN(
	[blob_in, scale, bias, blob_outs[1], blob_outs[2]], [blob_out],
	order=order, **kwargs)
	return blob_outputs
	else:
	blob_outputs = model.net.SpatialBN(
	[blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
	order=order, **kwargs)
	# Return the output
	return blob_outputs[0]


	def spatial_gn(model, blob_in, blob_out, dim_in,
	init_scale=1., init_bias=0.,
	ScaleInitializer=None, BiasInitializer=None,
	RunningMeanInitializer=None, RunningVarianceInitializer=None,
	order="NCHW", **kwargs):
	'''
	Group normalizes the input, cf. https://arxiv.org/abs/1803.08494.
	'''

	blob_out = blob_out or model.net.NextName()
	# Input: input, scale, bias
	# Output: output, group_mean, group_inv_std
	# scale: initialize with init_scale (default 1.)
	# [recommendation: set init_scale = 0. in the last layer for each res block]
	# bias: initialize with init_bias (default 0.)

	if model.init_params:
	scale_init = ("ConstantFill", {'value': init_scale})
	bias_init = ("ConstantFill", {'value': init_bias})

	ScaleInitializer = initializers.update_initializer(
	ScaleInitializer, scale_init, ("ConstantFill", {})
	)
	BiasInitializer = initializers.update_initializer(
	BiasInitializer, bias_init, ("ConstantFill", {})
	)
	else:
	ScaleInitializer = initializers.ExternalInitializer()
	BiasInitializer = initializers.ExternalInitializer()

	scale = model.create_param(
	param_name=blob_out + '_s',
	shape=[dim_in],
	initializer=ScaleInitializer,
	tags=ParameterTags.WEIGHT
	)

	bias = model.create_param(
	param_name=blob_out + '_b',
	shape=[dim_in],
	initializer=BiasInitializer,
	tags=ParameterTags.BIAS
	)

	blob_outs = [blob_out,
	blob_out + "_mean", blob_out + "_std"]

	blob_outputs = model.net.GroupNorm(
	[blob_in, scale, bias],
	blob_outs,
	**kwargs)
	# Return the output
	return blob_outputs[0]


	def layer_norm(
	model,
	blob_in,
	blob_out,
	dim_in,
	axis=1,
	epsilon=1e-4,
	initial_scale=1.0,
	initial_bias=0.0,
	):
	'''
	Layer normalizes the input, cf. https://arxiv.org/pdf/1607.06450.pdf.

	Args:
	blob_in: The input blob to layer normalize.
	blob_out: The layer normalized output blob.
	dim_in: The dimension of the scale and bias. For example, if blob_in is
	a 2D design matrix and axis is 1, this would be the number of
	columns.
	axis: (optional) The axis to normalize. Typically the feature axis.
	Defaults to 1.
	epsilon: (optional) A small value used for numerical stability in
	calculation. Defaults to 1e-4.
	initial_scale: (optional) The initial value for the learned scale
	parameter. Defaults to 1.0
	initial_bias: (optional) The initial value for the learned bias
	parameter of the layerwise standard deviation. Defaults to 0.0.

	Returns:
	A 3-tuple consisting of:
	- The layer normalized input blob.
	- The mean of the input blob across the given axis.
	- The standard deviation of the input blob acress the given axis.
	'''

	# The learned multiplicative scale or "gain".
	scale = model.create_param(
	param_name='{}_scale'.format(blob_out),
	shape=[dim_in] if isinstance(dim_in, int) else dim_in,
	initializer=initializers.Initializer(
	'ConstantFill',
	value=initial_scale,
	),
	tags=ParameterTags.WEIGHT,
	)

	# The learned additive bias or "shift".
	bias = model.create_param(
	param_name='{}_bias'.format(blob_out),
	shape=[dim_in] if isinstance(dim_in, int) else dim_in,
	initializer=initializers.Initializer(
	'ConstantFill',
	value=initial_bias,
	),
	tags=ParameterTags.BIAS,
	)

	normalized, mean, std = model.net.LayerNorm(
	[blob_in, scale, bias],
	[blob_out, blob_out + "_mean", blob_out + "_std"],
	axis=axis,
	epsilon=epsilon,
	elementwise_affine=True,
	)

	return normalized, mean, std

	def moments_with_running_stats(model, blob_in, blob_out, dim_in,
	RunningMeanInitializer=None, RunningVarianceInitializer=None,
	order="NCHW", **kwargs):

	if model.init_params:
	rm_init = ("ConstantFill", {'value': 0.0})
	riv_init = ("ConstantFill", {'value': 1.0})

	RunningMeanInitializer = initializers.update_initializer(
	RunningMeanInitializer, rm_init, ("ConstantFill", {})
	)
	RunningVarianceInitializer = initializers.update_initializer(
	RunningVarianceInitializer, riv_init, ("ConstantFill", {})
	)
	else:
	RunningMeanInitializer = initializers.ExternalInitializer()
	RunningVarianceInitializer = initializers.ExternalInitializer()

	running_mean = model.create_param(
	param_name=blob_out + '_rm',
	shape=[dim_in],
	initializer=RunningMeanInitializer,
	tags=ParameterTags.COMPUTED_PARAM
	)

	# this is just running variance
	running_inv_var = model.create_param(
	param_name=blob_out + '_riv',
	shape=[dim_in],
	initializer=RunningVarianceInitializer,
	tags=ParameterTags.COMPUTED_PARAM
	)

	blob_outs = [blob_out + "_sm", blob_out + "_sv"]
	if order == 'NCHW':
	blob_outputs = model.net.Moments(
	[blob_in], blob_outs,
	axes=[0, 2, 3],
	order=order, keepdims=False, **kwargs)
	elif order == 'NHWC':
	blob_outputs = model.net.Moments(
	[blob_in], blob_outs,
	axes=[0, 1, 2],
	order=order, keepdims=False, **kwargs)
	return blob_outputs