def add_update(dest: Tensor,
delta: Tensor,
*,
alpha: Union[Tensor, float, int] = 1.0,
beta: Union[Tensor, float, int] = 1.0,
bias: Union[Tensor, float, int] = 0.0):
r"""Inplace modify ``dest`` as follows:
.. math::
dest = alpha * dest + beta * delta + bias
:param dest: input data that will be inplace modified.
:param delta: update value that will be added to ``dest``.
:param alpha: weight ratio of ``dest``. Default: 1.0
:param beta: weight ratio of ``delta``. Default: 1.0
:param bias: bias value appended to the result. Default: 0.0
"""
if isinstance(beta, Tensor) or isinstance(alpha, Tensor):
delta *= beta
beta = 1.0
if isinstance(alpha, Tensor):
delta += (alpha - 1.0) * dest
alpha = 1.0
if isinstance(bias, Tensor):
delta += bias
bias = 0.0
comp_graph = dest._comp_graph or get_default_graph()
comp_node = dest._comp_node
if not isinstance(delta, Tensor):
_delta = mgb.make_immutable(value=delta,
comp_node=comp_node,
comp_graph=comp_graph)
else:
_delta = delta._attach(comp_graph)
_dest = dest._attach(comp_graph)
# use (dest, delta) as the key, so we could not add the same delta to dest in static graph
key = (comp_graph._id(), _dest.id, _delta.id)
if key in _add_update_cache:
_alpha, _beta, _bias, config = _add_update_cache[key]
mgb.mgb._mgb.SharedScalar__set(_alpha, alpha)
mgb.mgb._mgb.SharedScalar__set(_beta, beta)
mgb.mgb._mgb.SharedScalar__set(_bias, bias)
else:
_alpha = mgb.SharedScalar(alpha)
_beta = mgb.SharedScalar(beta)
_bias = mgb.SharedScalar(bias)
config = mgb.helper.gen_config(None, comp_node, None)
_add_update_cache[key] = (_alpha, _beta, _bias, config)
u = mgb.mgb._Opr.add_update(_dest, barrier(_delta), _alpha, _beta, _bias,
_dummy, config)
mark_impure(u)
return Tensor(u)
def scalar(
value,
dtype: type = None,
device: Optional[mgb.CompNode] = None,
comp_graph: Optional[mgb.CompGraph] = None,
) -> Tensor:
device, comp_graph = _use_default_if_none(device, comp_graph)
return Tensor(mgb.make_immutable(device, comp_graph, value, dtype=dtype, name=None))
def one_hot(inp: Tensor, num_classes: int) -> Tensor:
r"""
Perform one-hot encoding for the input tensor.
:param inp: input tensor
:param num_classes: number of classes denotes the last dimension of the output tensor
Examples:
.. testcode::
import numpy as np
from megengine import tensor
import megengine.functional as F
inp = tensor(np.arange(1, 4, dtype=np.int32))
out = F.one_hot(inp, num_classes=4)
print(out.numpy())
Outputs:
.. testoutput::
[[0 1 0 0]
[0 0 1 0]
[0 0 0 1]]
"""
comp_node, comp_graph = _decide_comp_node_and_comp_graph(inp)
zeros = mgb.make_immutable(value=0,
comp_node=comp_node,
comp_graph=comp_graph)
zeros_symvar = zeros.broadcast(inp.shapeof(), num_classes)
ones = mgb.make_immutable(value=1,
comp_node=comp_node,
comp_graph=comp_graph)
ones_symvar = ones.broadcast(inp.shapeof(), 1)
return Tensor(
mgb.opr.indexing_set_one_hot(zeros_symvar,
axis=len(inp.shapeof()),
index=inp,
value=ones_symvar))
def scalar(
value,
dtype: type = None,
device: Optional[mgb.CompNode] = None,
comp_graph: Optional[mgb.CompGraph] = None,
) -> Tensor:
"""
convert ``value`` to the type of :class:`~.Tensor`.
"""
device, comp_graph = _use_default_if_none(device, comp_graph)
return Tensor(
mgb.make_immutable(device, comp_graph, value, dtype=dtype, name=None))
def batch_norm2d(
inp: Tensor,
running_mean: Tensor,
running_var: Tensor,
weight: Optional[Tensor] = None,
bias: Optional[Tensor] = None,
training: bool = False,
momentum: float = 0.9,
eps: float = 1e-5,
) -> Tensor:
"""Applies batch normalization to the input.
:param inp: input tensor.
:param running_mean: tensor to store running mean.
:param running_var: tensor to store running variance.
:param weight: scaling tensor in the learnable affine parameters.
See :math:`\gamma` in :class:`~.BatchNorm2d`
:param bias: bias tensor in the learnable affine parameters.
See :math:`\beta` in :class:`~.BatchNorm2d`
:param training: a boolean value to indicate whether batch norm is performed
in traning mode. Default: ``False``
:param momentum: the value used for the ``running_mean`` and ``running_var``
computation.
Default: 0.9
:param eps: a value added to the denominator for numerical stability.
Default: 1e-5.
Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information.
"""
inp = mgb.opr.mark_no_broadcast_elemwise(inp)
_channels = inp.imm_shape[1]
_ndim = len(inp.imm_shape)
_param_shape = (1, _channels) + (1,) * (_ndim - 2)
assert _ndim == 4, "only 4D tensor supported"
if weight is not None:
weight = weight.reshape(*_param_shape)
else:
weight = mgb.make_immutable(*_use_default_if_none(None, None), 1.0).broadcast(
*_param_shape
)
if bias is not None:
bias = bias.reshape(*_param_shape)
else:
bias = mgb.make_immutable(*_use_default_if_none(None, None), 0.0).broadcast(
*_param_shape
)
FwdMode = mgb.opr_param_defs.BN.FwdMode
fwdmode = FwdMode.TRAINING if training else FwdMode.INFERENCE
avg_factor = 1 - momentum
if running_mean is not None and running_var is not None:
if training:
inp = barrier(inp)
output = mgb.opr.batch_norm(
inp,
weight,
bias,
running_mean,
running_var,
param_dim="DIM_1C11",
fwd_mode=fwdmode,
epsilon=eps,
avg_factor=avg_factor,
)[-1]
if training:
mark_impure(output)
else:
output = mgb.opr.batch_norm_no_statistic(
inp,
weight,
bias,
param_dim="DIM_1C11",
fwd_mode=fwdmode,
epsilon=eps,
avg_factor=avg_factor,
)[-1]
return output
def _dummy():
return mgb.make_immutable(*graph._use_default_if_none(None, None), 0)
def batch_norm2d(
inp: Tensor,
running_mean: Tensor,
running_var: Tensor,
weight: Optional[Tensor] = None,
bias: Optional[Tensor] = None,
training: bool = False,
momentum: float = 0.9,
eps: float = 1e-5,
) -> Tensor:
"""Applies batch normalization to the input.
:type inp: Tensor
:param inp: The input tensor.
:type num_features: int
:param num_features: usually the :math:`C` from an input of size
:math:`(N, C, H, W)` or the highest ranked dimension of an input with
less than 4D.
:type eps: float
:param eps: a value added to the denominator for numerical stability.
Default: 1e-5.
:type momentum: float
:param momentum: the value used for the `running_mean` and `running_var`
computation.
Default: 0.1
:type affine: bool
:param affine: a boolean value that when set to ``True``, this module has
learnable affine parameters. Default: ``True``
:type track_running_stats: bool
:param track_running_stats: when set to ``True``, this module tracks the
running mean and variance. When set to ``False``, this module does not
track such statistics and always uses batch statistics in both training
and eval modes. Default: ``True``.
Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information.
"""
inp = mgb.opr.mark_no_broadcast_elemwise(inp)
_channels = inp.imm_shape[1]
_ndim = len(inp.imm_shape)
_param_shape = (1, _channels) + (1,) * (_ndim - 2)
assert _ndim == 4, "only 4D tensor supported"
if weight is not None:
weight = weight.reshape(*_param_shape)
else:
weight = mgb.make_immutable(*_use_default_if_none(None, None), 1.0).broadcast(
*_param_shape
)
if bias is not None:
bias = bias.reshape(*_param_shape)
else:
bias = mgb.make_immutable(*_use_default_if_none(None, None), 0.0).broadcast(
*_param_shape
)
FwdMode = mgb.opr_param_defs.BN.FwdMode
fwdmode = FwdMode.TRAINING if training else FwdMode.INFERENCE
avg_factor = 1 - momentum
if running_mean is not None and running_var is not None:
if training:
inp = barrier(inp)
output = mgb.opr.batch_norm(
inp,
weight,
bias,
running_mean,
running_var,
param_dim="DIM_1C11",
fwd_mode=fwdmode,
epsilon=eps,
avg_factor=avg_factor,
)[-1]
if training:
mark_impure(output)
else:
output = mgb.opr.batch_norm_no_statistic(
inp,
weight,
bias,
param_dim="DIM_1C11",
fwd_mode=fwdmode,
epsilon=eps,
avg_factor=avg_factor,
)[-1]
return output