MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / batch_norm

Function batch_norm

imperative/python/megengine/functional/nn.py:1241–1345  ·  view source on GitHub ↗

r"""Applies batch normalization to the input. Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information. Args: inp: input tensor. running_mean: tensor to store running mean. running_var: tensor to store running variance. weight: scaling

(
    inp: Tensor,
    running_mean: Tensor = None,
    running_var: Tensor = None,
    weight: Optional[Tensor] = None,
    bias: Optional[Tensor] = None,
    *,
    training: bool = False,
    momentum: float = 0.9,
    eps: float = 1e-5,
    inplace: bool = True,
)

Source from the content-addressed store, hash-verified

1239
1240
1241def batch_norm(
1242 inp: Tensor,
1243 running_mean: Tensor = None,
1244 running_var: Tensor = None,
1245 weight: Optional[Tensor] = None,
1246 bias: Optional[Tensor] = None,
1247 *,
1248 training: bool = False,
1249 momentum: float = 0.9,
1250 eps: float = 1e-5,
1251 inplace: bool = True,
1252):
1253 r"""Applies batch normalization to the input.
1254
1255 Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information.
1256
1257 Args:
1258 inp: input tensor.
1259 running_mean: tensor to store running mean.
1260 running_var: tensor to store running variance.
1261 weight: scaling tensor in the learnable affine parameters.
1262 See :math:`\gamma` in :class:`~.BatchNorm2d`.
1263 bias: bias tensor in the learnable affine parameters.
1264 See :math:`\beta` in :class:`~.BatchNorm2d`.
1265 training: a boolean value to indicate whether batch norm is performed
1266 in training mode. Default: False
1267 momentum: value used for the ``running_mean`` and ``running_var``
1268 computation. Default: 0.9
1269 eps: a value added to the denominator for numerical stability. Default: 1e-5
1270 inplace: whether to update ``running_mean`` and ``running_var``
1271 inplace or return new tensors. Default: True
1272 compute_mode: When set to 'default', no special requirements will be
1273 placed on the precision of intermediate results. When set to 'float32',
1274 float32 would be used for accumulator and intermediate result, but only
1275 effective when input and output are of float16 dtype.
1276 param_dim: a value indicating in which format the parameters are.
1277 Default: 'dim_1c11', which means NCHW format.
1278 And 'dim_111c' means NHWC format.
1279 """
1280
1281 def make_full_if_none(x, value):
1282 x_ndim = None if x is None else x.ndim
1283 # in general case, x will be returned here directly
1284 if x_ndim is not None and x_ndim != 1:
1285 return x
1286
1287 C = inp.shape[1]
1288 pshape = (1, C, 1, 1)
1289
1290 if x is None:
1291 x = Const(value, inp.dtype, inp.device)
1292 shape = astensor1d(pshape, inp, dtype="int32", device=inp.device)
1293 (result,) = apply(builtin.Broadcast(), x, shape)
1294 result.format = inp.format
1295 return result
1296 else:
1297 assert x_ndim == 1
1298 shape = astensor1d(pshape, inp, dtype="int32", device=inp.device)

Callers 1

sync_batch_normFunction · 0.70

Calls 2

make_full_if_noneFunction · 0.85
applyFunction · 0.50

Tested by

no test coverage detected