提交 48822557 编写于 作者: M Megvii Engine Team

docs(mge/bn): fix docs and tests of batchnorm

GitOrigin-RevId: 8a96aa5fc221df1dbcf185b24d82cf7cf2a3cc24
上级 2c4d1afe
...@@ -126,7 +126,7 @@ class BatchNorm2d(_BatchNorm): ...@@ -126,7 +126,7 @@ class BatchNorm2d(_BatchNorm):
By default, during training this layer keeps running estimates of its By default, during training this layer keeps running estimates of its
computed mean and variance, which are then used for normalization during computed mean and variance, which are then used for normalization during
evaluation. The running estimates are kept with a default :attr:`momentum` evaluation. The running estimates are kept with a default :attr:`momentum`
of 0.1. of 0.9.
If :attr:`track_running_stats` is set to ``False``, this layer will not If :attr:`track_running_stats` is set to ``False``, this layer will not
keep running estimates, and batch statistics are instead used during keep running estimates, and batch statistics are instead used during
...@@ -154,7 +154,7 @@ class BatchNorm2d(_BatchNorm): ...@@ -154,7 +154,7 @@ class BatchNorm2d(_BatchNorm):
:type momentum: float :type momentum: float
:param momentum: the value used for the `running_mean` and `running_var` :param momentum: the value used for the `running_mean` and `running_var`
computation. computation.
Default: 0.1 Default: 0.9
:type affine: bool :type affine: bool
:param affine: a boolean value that when set to ``True``, this module has :param affine: a boolean value that when set to ``True``, this module has
learnable affine parameters. Default: ``True`` learnable affine parameters. Default: ``True``
...@@ -174,12 +174,18 @@ class BatchNorm2d(_BatchNorm): ...@@ -174,12 +174,18 @@ class BatchNorm2d(_BatchNorm):
# With Learnable Parameters # With Learnable Parameters
m = M.BatchNorm2d(4) m = M.BatchNorm2d(4)
inp = mge.tensor(np.random.rand(64, 4, 32, 32)) inp = mge.tensor(np.random.rand(1, 4, 3, 3).astype("float32"))
oup = m(inp) oup = m(inp)
print(m.weight, m.bias)
# Without Learnable Parameters # Without Learnable Parameters
m = M.BatchNorm2d(4, affine=False) m = M.BatchNorm2d(4, affine=False)
oup = m(inp) oup = m(inp)
print(m.weight, m.bias)
.. testoutput::
Tensor([1. 1. 1. 1.]) Tensor([0. 0. 0. 0.])
None None
""" """
def _check_input_ndim(self, inp): def _check_input_ndim(self, inp):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册