def _batch_norm(inputs, name, trainable=True, training=True):
params_shape = [inputs.shape[1]]
# Float32 required to avoid precision-loss when using fp16 input/output
params_dtype = flow.float32 if inputs.dtype == flow.float16 else inputs.dtype
if not flow.current_global_function_desc().IsTrainable() or not trainable:
training = False
with flow.scope.namespace(name):
beta = flow.get_variable(
name="beta",
shape=params_shape,
dtype=params_dtype,
initializer=flow.zeros_initializer(),
trainable=trainable,
distribute=distribute_util.broadcast(),
)
gamma = flow.get_variable(
name="gamma",
shape=params_shape,
dtype=params_dtype,
initializer=flow.ones_initializer(),
trainable=trainable,
distribute=distribute_util.broadcast(),
)
moving_mean = flow.get_variable(
name="moving_mean",
shape=params_shape,
dtype=params_dtype,
initializer=flow.zeros_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
)
moving_variance = flow.get_variable(
name="moving_variance",
shape=params_shape,
dtype=params_dtype,
initializer=flow.ones_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
)
builder = (flow.user_op_builder(
id_util.UniqueStr(name)).Op("normalization").Input(
"x", [inputs]).Input("moving_mean", [moving_mean]).Input(
"moving_variance",
[moving_variance]).Input("gamma", [gamma]).Input(
"beta", [beta]).Output("y").Attr("axis", 1).Attr(
"epsilon",
1.001e-5).Attr("training",
training).Attr("momentum", 0.997))
if trainable and training:
builder = builder.Output("mean").Output("inv_variance")
return builder.Build().InferAndTryRun().RemoteBlobList()[0]
def find_or_create_module(module_name, create, reuse=False):
assert callable(create)
sess = session_ctx.GetDefaultSession()
job_name = oneflow.current_global_function_desc().job_config_proto.job_name
if job_name not in sess.job_name2module_name2module_:
sess.job_name2module_name2module_[job_name] = {}
module_name2module = sess.job_name2module_name2module_[job_name]
if module_name not in module_name2module:
module = create()
assert isinstance(module, module_util.Module)
module_name2module[module_name] = module
else:
if not reuse:
assert module_name not in sess.existed_module_names_, (
"duplicated module_name `%s' in global_function `%s'" %
(module_name, job_name))
else:
# do nothing
pass
sess.existed_module_names_.add(module_name)
return module_name2module[module_name]
def Foo():
test_case.assertEqual(
flow.current_global_function_desc().IsTrainable(), False)
return flow.get_variable("w", (10, ),
initializer=flow.constant_initializer(1))
def batch_normalization(
inputs: remote_blob_util.BlobDef,
axis: int = -1,
momentum: float = 0.99,
epsilon: float = 0.001,
center: bool = True,
scale: bool = True,
beta_initializer: Optional[op_conf_util.InitializerConf] = None,
gamma_initializer: Optional[op_conf_util.InitializerConf] = None,
beta_regularizer: Optional[op_conf_util.RegularizerConf] = None,
gamma_regularizer: Optional[op_conf_util.RegularizerConf] = None,
moving_mean_initializer: Optional[op_conf_util.InitializerConf] = None,
moving_variance_initializer: Optional[op_conf_util.InitializerConf] = None,
trainable: bool = True,
training: bool = True,
name: str = "BatchNorm",
) -> remote_blob_util.BlobDef:
r"""Analogous to `tf.keras.layers.BatchNormalization <https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization>`_
Args:
inputs (remote_blob_util.BlobDef): Input `Blob`.
axis (int, optional): An int specifies the aixs that should be normalized . Default is -1, which normalizes the last axis.
momentum (float, optional): A float specifies the momontum for the moving average. Defaults to 0.99.
epsilon (float, optional): A small float added to avoid division by zero. Defaults to 0.001.
center (bool, optional): A boolean specifies whether to add offset to normalized `Blob`. Defaults to True.
scale (bool, optional): A boolean specifies whether to multiply normalized `Blob` by gamma. Defaults to True.
beta_initializer (Optional[op_conf_util.InitializerConf], optional): Initializer for beta. Defaults to None.
gamma_initializer (Optional[op_conf_util.InitializerConf], optional): Initializer for gamma. Defaults to None.
beta_regularizer (Optional[op_conf_util.RegularizerConf], optional): Regularizer for beta. Defaults to None.
gamma_regularizer (Optional[op_conf_util.RegularizerConf], optional): Regularizer for gamma. Defaults to None.
moving_mean_initializer (Optional[op_conf_util.InitializerConf], optional): Initializer for moving mean. Defaults to None.
moving_variance_initializer (Optional[op_conf_util.InitializerConf], optional): Initializer for moving variance. Defaults to None.
trainable (bool, optional): A boolean specifies whether to train variables. Defaults to True.
training (bool, optional): A boolean specifies whether now is training the model. Defaults to True.
name (Optional[str], optional): This layer's name. Defaults to None.
Returns:
remote_blob_util.BlobDef: A `Blob` with same shape of input.
Raises:
ValueError: If axis is out of dimension of input.
"""
if axis < 0:
axis += len(inputs.shape)
assert axis >= 0 and axis < len(inputs.shape)
params_shape = [inputs.shape[axis]]
# Float32 required to avoid precision-loss when using fp16 input/output
params_dtype = flow.float32 if inputs.dtype == flow.float16 else inputs.dtype
if not flow.current_global_function_desc().IsTrainable() or not trainable:
training = False
with flow.scope.namespace(name):
if center:
beta = flow.get_variable(
name="beta",
shape=params_shape,
dtype=params_dtype,
initializer=beta_initializer or flow.zeros_initializer(),
regularizer=beta_regularizer,
trainable=trainable,
distribute=distribute_util.broadcast(),
reuse=False,
)
else:
beta = flow.constant(0,
dtype=params_dtype,
shape=params_shape,
name="beta")
if scale:
gamma = flow.get_variable(
name="gamma",
shape=params_shape,
dtype=params_dtype,
initializer=gamma_initializer or flow.ones_initializer(),
regularizer=gamma_regularizer,
trainable=trainable,
distribute=distribute_util.broadcast(),
reuse=False,
)
else:
gamma = flow.constant(1,
dtype=params_dtype,
shape=params_shape,
name="gamma")
moving_mean = flow.get_variable(
name="moving_mean",
shape=params_shape,
dtype=params_dtype,
initializer=moving_mean_initializer or flow.zeros_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
reuse=False,
)
moving_variance = flow.get_variable(
name="moving_variance",
shape=params_shape,
dtype=params_dtype,
initializer=moving_variance_initializer or flow.ones_initializer(),
trainable=False,
distribute=distribute_util.broadcast(),
reuse=False,
)
if flow.current_scope().device_parallel_desc_symbol.device_tag == "cpu":
if training:
reduce_axis = []
for dim in range(len(inputs.shape)):
if dim != axis:
reduce_axis.append(dim)
mean, variance = flow.nn.moments(inputs,
reduce_axis,
keepdims=False)
def update_moving(moving, this_batch):
moving_identity = flow.identity(moving)
flow.assign(
moving,
momentum * moving_identity + (1 - momentum) * this_batch)
update_moving(moving_mean, mean)
update_moving(moving_variance, variance)
return flow.nn.batch_normalization(
x=inputs,
mean=mean,
variance=variance,
offset=beta,
scale=gamma,
variance_epsilon=epsilon,
axis=axis,
name=name,
)
else:
mean = moving_mean
variance = moving_variance
return flow.nn.batch_normalization(
x=inputs,
mean=mean,
variance=variance,
offset=beta,
scale=gamma,
variance_epsilon=epsilon,
axis=axis,
name=name,
)
else:
builder = (flow.user_op_builder(name).Op("normalization").Input(
"x", [inputs]).Input("moving_mean", [moving_mean]).Input(
"moving_variance",
[moving_variance]).Input("gamma", [gamma]).Input(
"beta", [beta]).Output("y").Attr("axis", axis).Attr(
"epsilon",
epsilon).Attr("training",
training).Attr("momentum", momentum))
if trainable and training:
builder = builder.Output("mean").Output("inv_variance")
return builder.Build().InferAndTryRun().RemoteBlobList()[0]
def moving_average_min_max_observer(
input: oneflow_api.BlobDesc,
quantization_bit: int = 8,
quantization_scheme: str = "symmetric",
quantization_formula: str = "google",
momentum: float = 0.95,
name: Optional[str] = None,
) -> Tuple[oneflow_api.BlobDesc, oneflow_api.BlobDesc]:
r"""Compute the quantization parameters based on the moving average of the input tensor's min and max values.
First compute the moving\_max and moving\_min value of input tensor:
if quantization_scheme == "symmetric":
.. math::
& moving\_max = moving\_max * momentum + |max(input)| * (1 - momentum)
& moving\_min = moving\_max
elif quantization_scheme == "affine":
.. math::
& moving\_max = moving\_max * momentum + max(input) * (1 - momentum)
& moving\_min = moving\_min * momentum + min(input) * (1 - momentum)
The moving average of min and max values are initialized as the first batch of input `Blob`'s min and max.
Then compute the scale and zero_point with the following equations:
if quantization_scheme == "symmetric":
.. math::
& denom = 2^{quantization\_to\_bit - 1} - 1
& scale = moving\_max / denom
& zero\_point = 0
elif quantization_scheme == "affine":
.. math::
& denom = 2^{quantization\_to\_bit} - 1
& scale = (moving\_max - moving\_min) / denom
& zero\_point = -moving\_min / scale
Args:
input (oneflow_api.BlobDesc): input tensor.
quantization_bit (int): Quantize input to uintX / intX, X can be in range [2, 8]. Defaults to 8.
quantization_scheme (str): "symmetric" or "affine", quantize to signed / unsigned integer. Defaults to "symmetric".
quantization_formula (str): Support "google" or "cambricon".
momentum (float): Smoothing parameter for exponential moving average operation. Defaults to 0.95.
name (Optional[str]): This operator's name. Defaults to None.
Returns:
Tuple[oneflow_api.BlobDesc, oneflow_api.BlobDesc]: The scale and zero_point of input tensor.
For example:
.. code-block:: python
import oneflow as flow
import numpy as np
import oneflow.typing as tp
@flow.global_function(type="predict", function_config=flow.FunctionConfig())
def QuantizeJob(
input: tp.Numpy.Placeholder(input_shape, dtype=type_name_to_flow_type[dtype])
): tp.Numpy
with flow.scope.placement(device_type, "0:0"):
scale, zero_point = flow.quantization.moving_average_min_max_observer(
input, quantization_bit=8,
quantization_scheme="symmetric",
quantization_formula="google",
momentum=0.95
)
return scale, zero_point
input = (np.random.random(input_shape) - 0.5).astype(type_name_to_np_type[dtype])
scale, zero_point = QuantizeJob(input)
"""
op_name = (name if name is not None else
id_util.UniqueStr("MovingAverageMinMaxObserver_"))
training = True if flow.current_global_function_desc().IsTrainable(
) else False
with flow.scope.namespace(op_name):
moving_max = flow.get_variable(
"moving_max",
shape=(1, ),
dtype=input.dtype,
initializer=flow.zeros_initializer(input.dtype),
trainable=False,
)
moving_min = flow.get_variable(
"moving_min",
shape=(1, ),
dtype=input.dtype,
initializer=flow.zeros_initializer(input.dtype),
trainable=False,
)
current_train_step = flow.get_variable(
"current_train_step",
shape=(1, ),
dtype=flow.int64,
initializer=flow.zeros_initializer(flow.int64),
trainable=False,
)
stop_update_after_iters = 1
scale, zero_point = (flow.user_op_builder(op_name).Op(
"moving_average_min_max_observer").Input("in", [input]).Input(
"current_train_step",
[current_train_step]).Input("moving_max", [moving_max]).Input(
"moving_min",
[moving_min]).Output("scale").Output("zero_point").Attr(
"training",
training).Attr("stop_update_after_iters",
stop_update_after_iters).Attr(
"quantization_bit",
quantization_bit).Attr(
"quantization_scheme",
quantization_scheme).Attr(
"quantization_formula",
quantization_formula).Attr(
"momentum", momentum).
Build().InferAndTryRun().RemoteBlobList())
return scale, zero_point
def global_mode_current_scope():
job_name = oneflow.current_global_function_desc().job_config_proto.job_name
return session_ctx.GetDefaultSession().GetCurrentScope(job_name)
