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 prelu(
inputs: remote_blob_util.BlobDef,
alpha_initializer: Optional[op_conf_util.InitializerConf] = None,
alpha_regularizer: Optional[op_conf_util.RegularizerConf] = None,
shared_axes: Optional[Sequence[int]] = None,
trainable: bool = True,
name: str = "PRelu",
model_distribute: distribute_util.Distribute = distribute_util.broadcast(),
) -> remote_blob_util.BlobDef:
alpha_shape = list(inputs.shape[1:])
if shared_axes is not None:
for i in shared_axes:
assert i >= 1 and i < len(inputs.shape)
alpha_shape[i - 1] = 1
if alpha_initializer is None:
alpha_initializer = flow.constant_initializer(0)
with flow.scope.namespace(name):
alpha = flow.get_variable(
name="alpha",
shape=alpha_shape,
dtype=inputs.dtype,
initializer=alpha_initializer,
regularizer=alpha_regularizer,
trainable=trainable,
distribute=model_distribute,
reuse=False,
)
op = (flow.user_op_builder(name).Op("prelu").Input("x", [inputs]).Input(
"alpha", [alpha]).Output("y").Build())
return op.InferAndTryRun().SoleOutputBlob()
def get_eager_variable(
name,
shape=None,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
distribute=distribute_util.broadcast(),
reuse=True,
):
assert isinstance(name, str)
assert isinstance(
shape, (list, tuple)
), "param shape should be a list or tuple of dimension"
job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName()
name = name_scope.GetJobNameScopePrefix(job_name) + name
sess = session_ctx.GetDefaultSession()
var_blob, job_var_blob = sess.TryGetVariableBlobOfJobFromStash(job_name, name)
if reuse is False:
assert job_var_blob is None, (
"varaible '{}' already exists, "
"getting the same variable is not allowed "
"when reuse is False".format(name)
)
if job_var_blob is None:
op_conf = _GenerateVariableOpConf(
name=name,
shape=shape,
dtype=dtype,
initializer=initializer,
regularizer=regularizer,
trainable=trainable,
model_name=model_name,
random_seed=random_seed,
distribute=distribute,
)
op_attribute = compile_context.CurJobAddConsistentOp(op_conf)
if var_blob is None:
var_blob = _CreateEagerVariableBlob(op_attribute)
op_executor.EagerInitVariableBlob(sess, op_conf, var_blob)
assert isinstance(var_blob, remote_blob_util.EagerConsistentBlob)
sess.StashVariableBlob4Job(job_name, op_conf.name, var_blob)
else:
assert isinstance(job_var_blob, remote_blob_util.EagerConsistentBlob)
assert isinstance(var_blob, remote_blob_util.EagerConsistentBlob)
assert var_blob.IdenticalTo(job_var_blob)
bw_blob_register = gradient_util.GetDefaultBackwardBlobRegister()
bw_blob_register.TrySetObject4BlobName(
var_blob.logical_blob_name, var_blob.blob_object
)
return var_blob
def GenerateVariableOpConf(
name,
shape,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
distribute=distribute_util.broadcast(),
):
op_conf = op_conf_util.OperatorConf()
op_conf.name = name
op_conf.variable_conf.shape.dim.extend(shape)
assert dtype is not None
op_conf.variable_conf.data_type = dtype.oneflow_proto_dtype
if rt_mode.CurrentMode() == rt_mode.NORMAL_MODE:
root_path = None
else:
root_path = (
compile_context.GetCurJobConfigProto().default_initialize_with_snapshot_path()
)
dir_path = os.path.join(root_path, name)
file_path = os.path.join(dir_path, "out")
if root_path and os.path.isfile(file_path):
op_conf.variable_conf.initialize_with_snapshot.path = dir_path
op_conf.variable_conf.initialize_with_snapshot.key = "out"
else:
if root_path:
print("{} not found, will be initialized".format(file_path))
if initializer is not None:
op_conf.variable_conf.initializer.CopyFrom(initializer)
if regularizer is not None:
op_conf.variable_conf.regularizer.CopyFrom(regularizer)
if trainable is not None:
op_conf.trainable = trainable
if model_name is not None:
op_conf.variable_conf.model_name = model_name
if type(distribute) is distribute_util.SplitDistribute:
op_conf.variable_conf.split_axis.value = distribute.axis
else:
op_conf.variable_conf.split_axis.ClearField("value")
if random_seed is not None:
op_conf.variable_conf.random_seed = random_seed
op_conf.variable_conf.out = "out"
return op_conf
def get_lazy_variable(
name,
shape=None,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
distribute=distribute_util.broadcast(),
reuse=True,
):
assert isinstance(name, str)
assert isinstance(
shape,
(list, tuple)), "param shape should be a list or tuple of dimension"
job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName()
name = name_scope.GetJobNameScopePrefix(job_name) + name
sess = session_ctx.GetDefaultSession()
var_blob, job_var_blob = sess.TryGetVariableBlobOfJobFromStash(
job_name, name)
if reuse is False:
assert job_var_blob is None, (
"varaible '{}' already exists, "
"getting the same variable is not allowed "
"when param reuse is False".format(name))
if job_var_blob is None:
op_conf = GenerateVariableOpConf(
name=name,
shape=shape,
dtype=dtype,
initializer=initializer,
regularizer=regularizer,
trainable=trainable,
model_name=model_name,
random_seed=random_seed,
distribute=distribute,
)
job_var_blob = _CreateVariableBlob(op_conf)
assert isinstance(job_var_blob, remote_blob_util.LazyConsistentBlob)
sess.StashVariableBlob4Job(job_name, op_conf.name, job_var_blob)
if var_blob is not None:
assert isinstance(var_blob, remote_blob_util.LazyConsistentBlob)
assert var_blob.IdenticalTo(job_var_blob)
else:
assert isinstance(job_var_blob, remote_blob_util.LazyConsistentBlob)
assert isinstance(var_blob, remote_blob_util.LazyConsistentBlob)
assert var_blob.IdenticalTo(job_var_blob)
return job_var_blob
def api_get_variable(
name: str,
shape: Optional[Sequence[int]] = None,
dtype: Optional[dtype_util.dtype] = dtype_util.float32,
initializer: Optional[op_conf_util.InitializerConf] = None,
regularizer: Optional[op_conf_util.RegularizerConf] = None,
trainable: Optional[bool] = None,
model_name: Optional[str] = None,
random_seed: Optional[int] = None,
distribute: distribute_util.Distribute = distribute_util.broadcast(),
reuse: bool = True,
) -> remote_blob_util.BlobDef:
r"""Create a variable or retrieve an existing one.
Args:
name: Name of this variable. One variable could be shared by multiple OneFlow functions. `None` by defauilt
shape: Shape of the variable. `None` by defauilt
dtype: Data type of the variable. `None` by defauilt
initializer: A initializer object. For instance, a :func:`~oneflow.ones_initializer`. `None` by defauilt
trainable: A `bool` to indicate if this variable is trainable. `True` by defauilt
model_name: A `string`. `'weight'` or `'bias'`. `None` by defauilt
random_seed: Random seed for random initializers. `None` by defauilt
"""
api = enable_if.unique([get_lazy_variable, get_eager_variable])
return api(
name,
shape=shape,
dtype=dtype,
initializer=initializer,
regularizer=regularizer,
trainable=trainable,
model_name=model_name,
random_seed=random_seed,
distribute=distribute,
reuse=reuse,
)
def api_get_variable(
name: str,
shape: Optional[Sequence[int]] = None,
dtype: Optional[dtype_util.dtype] = dtype_util.float32,
initializer: Optional[op_conf_util.InitializerConf] = None,
regularizer: Optional[op_conf_util.RegularizerConf] = None,
trainable: Optional[bool] = None,
model_name: Optional[str] = None,
random_seed: Optional[int] = None,
distribute: distribute_util.Distribute = distribute_util.broadcast(),
reuse: bool = True,
) -> remote_blob_util.BlobDef:
r"""Create a variable or retrieve an existing one.
Args:
name: Name of this variable. One variable could be shared by multiple OneFlow functions. `None` by defauilt
shape: Shape of the variable. `None` by defauilt
dtype: Data type of the variable. `None` by defauilt
initializer: A initializer object. For instance, a :func:`~oneflow.ones_initializer`. `None` by defauilt
trainable: A `bool` to indicate if this variable is trainable. `True` by defauilt
model_name: A `string`. `'weight'` or `'bias'`. `None` by defauilt
random_seed: Random seed for random initializers. `None` by defauilt
For example:
Example 1:
.. code-block:: python
import oneflow as flow
import oneflow.typing as tp
def watch_handler(y: tp.Numpy):
print("out", y)
@flow.global_function()
def variable_Job() -> None:
init = flow.constant_initializer(1.25)
variable = flow.get_variable(
"variable-weight",
shape=(1, 3, 2, 2),
initializer=init,
trainable=True
)
flow.watch(variable, watch_handler)
checkpoint = flow.train.CheckPoint()
checkpoint.init()
variable_Job()
# out [[[[1.25 1.25]
# [1.25 1.25]]
# [[1.25 1.25]
# [1.25 1.25]]
# [[1.25 1.25]
# [1.25 1.25]]]]
Example 2:
.. code-block:: python
import oneflow as flow
import numpy as np
import oneflow.typing as tp
def conv2d(input, filters, kernel_size, strides, padding, name):
input_shape = input.shape
weight_initializer = flow.truncated_normal(0.1)
weight_regularizer = flow.regularizers.l2(0.0005)
weight_shape = (filters,
input_shape[1],
kernel_size[0],
kernel_size[1])
weight = flow.get_variable(
name + "-weight",
shape=weight_shape,
initializer=weight_initializer,
regularizer=weight_regularizer,
)
return flow.nn.conv2d(input, weight, strides, padding, name=name)
@flow.global_function()
def conv2d_Job(x: tp.Numpy.Placeholder((1, 64, 32, 32))
) -> tp.Numpy:
conv = conv2d(x,
filters=128,
kernel_size=[3, 3],
strides=2,
padding='SAME',
name="Convlayer")
return conv
x = np.random.randn(1, 64, 32, 32).astype(np.float32)
out = conv2d_Job(x)
# out.shape (1, 128, 16, 16)
"""
api = enable_if.unique([get_lazy_variable, get_eager_variable])
return api(
name,
shape=shape,
dtype=dtype,
initializer=initializer,
regularizer=regularizer,
trainable=trainable,
model_name=model_name,
random_seed=random_seed,
distribute=distribute,
reuse=reuse,
)
def with_broadcast_distribute(self):
return self.with_distribute(distribute_util.broadcast())
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 layer_norm(
inputs: remote_blob_util.BlobDef,
center: bool = True,
scale: bool = True,
trainable: bool = True,
begin_norm_axis: int = 1,
begin_params_axis: int = -1,
epsilon: float = 1e-5,
name: str = "LayerNorm",
) -> remote_blob_util.BlobDef:
r"""Analogous to `tf.keras.layers.LayerNormalization <https://www.tensorflow.org/api_docs/python/tf/keras/layers/LayerNormalization>`_
Args:
inputs (remote_blob_util.BlobDef): Input `Blob`.
center (bool, optional): A boolean specifies whether to shift input `Blob`. Defaults to True.
scale (bool, optional): A boolean specifies whether to scaleinput `Blob`. Defaults to True.
trainable (bool, optional): A boolean specifies whether to train variables. Defaults to True.
begin_norm_axis (int, optional): An integer specifies which axis to normalize at first. Defaults to 1.
begin_params_axis (int, optional): An integer specifies which axis params at . Defaults to -1.
epsilon (float, optional): A small float is added to avoid division by zero. Defaults to 1e-5.
name (Optional[str], optional): This layer's name. Defaults to None.
Returns:
remote_blob_util.BlobDef: A normalized `Blob` with same shape of input.
"""
if center is False and scale is False:
trainable = False
beta = None
gamma = None
param_shape = inputs.shape[begin_params_axis:]
if center:
with flow.scope.namespace(name):
beta = flow.get_variable(
name="beta",
shape=param_shape,
dtype=inputs.dtype,
initializer=flow.constant_initializer(0.0),
trainable=trainable,
model_name="beta",
distribute=distribute_util.broadcast(),
reuse=False,
)
if scale:
with flow.scope.namespace(name):
gamma = flow.get_variable(
name="gamma",
shape=param_shape,
dtype=inputs.dtype,
initializer=flow.constant_initializer(1.0),
trainable=trainable,
model_name="gamma",
distribute=distribute_util.broadcast(),
reuse=False,
)
if flow.current_scope().device_parallel_desc_symbol.device_tag == "cpu":
if begin_norm_axis < 0:
begin_norm_axis = begin_norm_axis + len(inputs.shape)
reduce_axis = []
for dim in range(len(inputs.shape)):
if dim >= begin_norm_axis:
reduce_axis.append(dim)
mean, variance = flow.nn.moments(inputs, reduce_axis, keepdims=True)
axis = begin_norm_axis
normalized = flow.nn.batch_normalization(
x=inputs,
mean=mean,
variance=variance,
variance_epsilon=epsilon,
axis=axis,
name=name,
)
nd_params_shape = [1] * (len(inputs.shape) -
len(param_shape)) + list(param_shape)
affined = normalized
if gamma:
gamma = flow.reshape(gamma, nd_params_shape)
affined *= scale
if beta:
beta = flow.reshape(beta, nd_params_shape)
affined += beta
return affined
elif flow.current_scope().device_parallel_desc_symbol.device_tag == "gpu":
op_builder = (flow.user_op_builder(name).Op("layer_norm").Input(
"x", [inputs]).Output("y").Output("mean").Output("inv_variance"))
if beta is not None:
op_builder.Input("beta", [beta])
if gamma is not None:
op_builder.Input("gamma", [gamma])
op_builder.Output("normalized")
op_builder.Attr("center", center)
op_builder.Attr("scale", scale)
op_builder.Attr("begin_norm_axis", begin_norm_axis)
op_builder.Attr("begin_params_axis", begin_params_axis)
op_builder.Attr("epsilon", epsilon)
return op_builder.Build().InferAndTryRun().RemoteBlobList()[0]
else:
raise NotImplementedError
def dense(
inputs: remote_blob_util.BlobDef,
units: int,
activation: Optional[Callable[[remote_blob_util.BlobDef, str],
remote_blob_util.BlobDef]] = None,
use_bias: bool = True,
kernel_initializer: Optional[op_conf_util.InitializerConf] = None,
bias_initializer: Optional[op_conf_util.InitializerConf] = None,
kernel_regularizer: Optional[op_conf_util.RegularizerConf] = None,
bias_regularizer: Optional[op_conf_util.RegularizerConf] = None,
trainable: bool = True,
name: str = "Dense",
model_distribute: distribute_util.Distribute = distribute_util.broadcast(),
) -> remote_blob_util.BlobDef:
r"""Analogous to `tf.keras.layers.Dense <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense>`_
Args:
inputs (remote_blob_util.BlobDef): A 2D input `Blob`.
units (int): A positive integer for the dimensionality of the output space.
activation (Optional[remote_blob_util.BlobDef], optional): Activation function. Defaults to None.
use_bias (bool, optional): A boolean specifies whether to use a bias vector. Defaults to True.
kernel_initializer (Optional[op_conf_util.InitializerConf], optional): Initializer for the kernel weights matrix. Defaults to None.
bias_initializer (Optional[op_conf_util.InitializerConf], optional): [description]. Defaults to None.
kernel_regularizer (Optional[op_conf_util.RegularizerConf], optional): [description]. Defaults to None.
bias_regularizer (Optional[op_conf_util.RegularizerConf], optional): Regularizer for the bias vector. Defaults to None.
trainable (bool, optional): A boolean specifies whether to train the variables. Defaults to True.
name (Optional[str], optional): This layer's name. Defaults to None.
model_distribute (distribute_util.Distribute, optional): Define the way to ditribute the model. Defaults to distribute_util.broadcast().
Returns:
remote_blob_util.BlobDef: A N-D `Blob` with the shape of (batch_size, units).
Raises:
ValueError: The dimension of input `Blob` must be less than 2.
VauleError: Model distribute must be in auto, broadcast, split.
ValueError: The input must be a 2D `Blob` when the model distribute is split.
"""
in_shape = inputs.shape
in_num_axes = len(in_shape)
assert in_num_axes >= 2
assert (model_distribute is distribute_util.auto()
or model_distribute is distribute_util.broadcast()
or model_distribute is distribute_util.split(0))
if model_distribute is distribute_util.split(0):
assert in_num_axes == 2 # model distribute is hard for reshape split dim 1
if in_num_axes > 2:
inputs = flow.reshape(inputs, (-1, in_shape[-1]))
with flow.scope.namespace(name):
if kernel_initializer is None:
kernel_initializer = flow.constant_initializer(0)
weight = flow.get_variable(
name="weight",
shape=(units, inputs.shape[1]),
dtype=inputs.dtype,
initializer=kernel_initializer,
regularizer=kernel_regularizer,
trainable=trainable,
model_name="weight",
distribute=model_distribute,
reuse=False,
)
weight = weight.with_distribute(model_distribute)
out = flow.matmul(a=inputs, b=weight, transpose_b=True, name="matmul")
if use_bias:
if bias_initializer is None:
bias_initializer = flow.constant_initializer(0)
bias = flow.get_variable(
name="bias",
shape=(units, ),
dtype=inputs.dtype,
initializer=bias_initializer,
regularizer=bias_regularizer,
trainable=trainable,
model_name="bias",
distribute=model_distribute,
reuse=False,
)
bias = bias.with_distribute(model_distribute)
out = flow.nn.bias_add(out, bias, name="bias_add")
if callable(activation):
out = activation(out, name="activation")
if in_num_axes > 2:
out = flow.reshape(out, in_shape[:-1] + (units, ))
return out
def layer_norm(
inputs: remote_blob_util.BlobDef,
center: bool = True,
scale: bool = True,
trainable: bool = True,
begin_norm_axis: int = 1,
begin_params_axis: int = -1,
epsilon: float = 1e-5,
name: str = "LayerNorm",
) -> remote_blob_util.BlobDef:
r"""Analogous to `tf.keras.layers.LayerNormalization <https://www.tensorflow.org/api_docs/python/tf/keras/layers/LayerNormalization>`_
Args:
inputs (remote_blob_util.BlobDef): Input `Blob`.
center (bool, optional): A boolean specifies whether to shift input `Blob`. Defaults to True.
scale (bool, optional): A boolean specifies whether to scaleinput `Blob`. Defaults to True.
trainable (bool, optional): A boolean specifies whether to train variables. Defaults to True.
begin_norm_axis (int, optional): An integer specifies which axis to normalize at first. Defaults to 1.
begin_params_axis (int, optional): An integer specifies which axis params at . Defaults to -1.
epsilon (float, optional): A small float is added to avoid division by zero. Defaults to 1e-5.
name (Optional[str], optional): This layer's name. Defaults to None.
Returns:
remote_blob_util.BlobDef: A normalized `Blob` with same shape of input.
"""
op_builder = (
flow.user_op_builder(name)
.Op("layer_norm")
.Input("x", [inputs])
.Output("y")
.Output("mean")
.Output("inv_variance")
)
if center is False and scale is False:
trainable = False
param_shape = inputs.shape[begin_params_axis:]
if center:
with flow.scope.namespace(name):
beta = flow.get_variable(
name="beta",
shape=param_shape,
dtype=inputs.dtype,
initializer=flow.constant_initializer(0.0),
trainable=trainable,
model_name="beta",
distribute=distribute_util.broadcast(),
reuse=False,
)
op_builder.Input("beta", [beta])
if scale:
with flow.scope.namespace(name):
gamma = flow.get_variable(
name="gamma",
shape=param_shape,
dtype=inputs.dtype,
initializer=flow.constant_initializer(1.0),
trainable=trainable,
model_name="gamma",
distribute=distribute_util.broadcast(),
reuse=False,
)
op_builder.Input("gamma", [gamma])
op_builder.Output("normalized")
op_builder.Attr("center", center)
op_builder.Attr("scale", scale)
op_builder.Attr("begin_norm_axis", begin_norm_axis)
op_builder.Attr("begin_params_axis", begin_params_axis)
op_builder.Attr("epsilon", epsilon)
return op_builder.Build().InferAndTryRun().RemoteBlobList()[0]
def prelu(
inputs: remote_blob_util.BlobDef,
alpha_initializer: Optional[op_conf_util.InitializerConf] = None,
alpha_regularizer: Optional[op_conf_util.RegularizerConf] = None,
shared_axes: Optional[Sequence[int]] = None,
trainable: bool = True,
name: str = "PRelu",
model_distribute: distribute_util.Distribute = distribute_util.broadcast(),
) -> remote_blob_util.BlobDef:
r"""The Prelu(Parametric Rectified Linear Unit) activation.
The :math:`\alpha` is a parameter that can be trained in network
The equation is
.. math::
out = max(0, x) + \alpha*min(0, x)
Args:
inputs (remote_blob_util.BlobDef): The input Blob.
alpha_initializer (Optional[op_conf_util.InitializerConf], optional): The initializer of alpha. Defaults to None.
alpha_regularizer (Optional[op_conf_util.RegularizerConf], optional): The regularizer of alpha. Defaults to None.
shared_axes (Optional[Sequence[int]], optional): The axis along which to share learnable parameters for the prelu activation function. Defaults to None.
trainable (bool, optional): Whether to train the parameter :math:`\alpha`. Defaults to True.
name (str, optional): The name for the operation. Defaults to "PRelu".
model_distribute (distribute_util.Distribute, optional): Define the way to ditribute the model. Defaults to distribute_util.broadcast().
Returns:
remote_blob_util.BlobDef: The activated Blob
For example:
.. code-block:: python
import oneflow as flow
import oneflow.typing as tp
BATCH_SIZE = 100
def lenet(data, train=False):
initializer = flow.truncated_normal(0.1)
conv1 = flow.layers.conv2d(
data,
32,
5,
padding="SAME",
name="conv1",
kernel_initializer=initializer,
)
prelu1 = flow.layers.prelu(conv1,
alpha_initializer=initializer,
shared_axes=[2, 3],
name="Prelu1")
pool1 = flow.nn.max_pool2d(
prelu1, ksize=2, strides=2, padding="SAME", name="pool1", data_format="NCHW"
)
conv2 = flow.layers.conv2d(
pool1,
64,
5,
padding="SAME",
name="conv2",
kernel_initializer=initializer,
)
prelu2 = flow.layers.prelu(conv2,
alpha_initializer=initializer,
shared_axes=[2, 3],
name="Prelu2")
pool2 = flow.nn.max_pool2d(
prelu2, ksize=2, strides=2, padding="SAME", name="pool2", data_format="NCHW"
)
reshape = flow.reshape(pool2, [pool2.shape[0], -1])
hidden = flow.layers.dense(
reshape,
512,
activation=flow.nn.relu,
kernel_initializer=initializer,
name="dense1",
)
if train:
hidden = flow.nn.dropout(hidden, rate=0.5, name="dropout")
return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name="dense2")
@flow.global_function(type="train")
def train_job(
images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
with flow.scope.placement("gpu", "0:0"):
logits = lenet(images, train=True)
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
labels, logits, name="softmax_loss"
)
lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
flow.optimizer.SGD(lr_scheduler, momentum=0.9).minimize(loss)
return loss
"""
alpha_shape = list(inputs.shape[1:])
if shared_axes is not None:
for i in shared_axes:
assert i >= 1 and i < len(inputs.shape)
alpha_shape[i - 1] = 1
if alpha_initializer is None:
alpha_initializer = flow.constant_initializer(0)
with flow.scope.namespace(name):
alpha = flow.get_variable(
name="alpha",
shape=alpha_shape,
dtype=inputs.dtype,
initializer=alpha_initializer,
regularizer=alpha_regularizer,
trainable=trainable,
distribute=model_distribute,
reuse=False,
)
op = (flow.user_op_builder(name).Op("prelu").Input("x", [inputs]).Input(
"alpha", [alpha]).Output("y").Build())
return op.InferAndTryRun().SoleOutputBlob()
