Function test_instance_norm

imperative/python/test/unit/module/test_module.py:794–859

(affine)

@pytest.mark.parametrize("affine", [True, False])
def test_instance_norm(affine):
    num_channels = 4
    weight_np = np.random.uniform(-0.5, 0.5, (num_channels))
    bias_np = np.random.uniform(-0.5, 0.5, (num_channels))

    class OriginInstanceNormFunc(Module):
        def __init__(self, eps=1e-5, affine=True, **kwargs):
            super().__init__(**kwargs)
            self.num_channels = num_channels
            self.eps = eps
            self.affine = affine
            if self.affine:
                self.weight = Parameter(weight_np)
                self.bias = Parameter(bias_np)
            else:
                self.weight = None
                self.bias = None

        def forward(self, x):
            N, C, H, W = x.shape
            x = x.reshape(N, self.num_channels, -1)
            mean = x.mean(axis=2, keepdims=True)
            var = (x * x).mean(axis=2, keepdims=True) - mean * mean
            x = (x - mean) / F.sqrt(var + self.eps)
            x = x.reshape(N, C, H, W)
            if self.affine:
                x = self.weight.reshape(1, -1, 1, 1) * x + self.bias.reshape(
                    1, -1, 1, 1
                )
            return x

    inp = np.random.uniform(-0.5, 0.5, (2, num_channels, 10, 16)).astype("float32")
    mge_inp = Tensor(inp)
    mge_m = InstanceNorm(num_channels, affine=affine)
    mge_m.weight = Parameter(weight_np)
    mge_m.bias = Parameter(bias_np)

    ori_inp = Tensor(inp)
    ori_m = OriginInstanceNormFunc(affine=affine)

    mge_im = mge.autodiff.GradManager().attach((*mge_m.parameters(), mge_inp))
    ori_im = mge.autodiff.GradManager().attach((*ori_m.parameters(), ori_inp))
    dy = Tensor(np.random.uniform(-0.5, 0.5, inp.shape))

    for i in range(2):
        with mge_im:
            mge_output = mge_m(mge_inp)

            mge_im.backward(mge_output, dy)

        with ori_im:
            ori_output = ori_m(ori_inp)

            ori_im.backward(ori_output, dy)

        np.testing.assert_allclose(mge_output.numpy(), ori_output.numpy(), atol=1e-05)
        np.testing.assert_allclose(
            ori_inp.grad.numpy(), mge_inp.grad.numpy(), atol=1e-04