Got different result using AdaptiveInstanceNorm2d

❓Question

I want to use AdaptiveInstanceNorm2d layer to build my model, but when I convert my pytorch model to coreml model, AdaptiveInstanceNorm2d layer got different output than pytorch model

to reproduce:

import torch
from torch import nn
import torch.nn.functional as F
import coremltools

class AdaptiveInstanceNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        # weight and bias are dynamically assigned
        self.weight = torch.rand(num_features)
        self.bias = torch.rand(num_features)
        # just dummy buffers, not used
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, x):
        assert self.weight is not None and self.bias is not None, "Please assign weight and bias before calling AdaIN!"
        b, c = x.size(0), x.size(1)
        running_mean = self.running_mean.repeat(b).type_as(x)
        running_var = self.running_var.repeat(b).type_as(x)

        # Apply instance norm
        x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])

        out = F.batch_norm(
            x_reshaped, running_mean, running_var, self.weight, self.bias,
            True, self.momentum, self.eps)

        return out.view(b, c, *x.size()[2:])

    def __repr__(self):
        return self.__class__.__name__ + '(' + str(self.num_features) + ')'

class TestModule(nn.Module):
    def __init__(self):
        super(TestModule, self).__init__()
        self.norm = AdaptiveInstanceNorm2d(3)
    def forward(self, x):
        return self.norm(x)
x = torch.rand(1,3,100,100)
test_module = TestModule()
pytorch_output = test_module(x).detach().numpy()

from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
@register_torch_op(override=True)
def type_as(context, node):
    inputs = _get_inputs(context, node)
    context.add(mb.cast(x=inputs[0], dtype='fp32'), node.name)

traced_test_module= torch.jit.trace(test_module, x)
test_module_mlmodel = coremltools.convert(
    traced_test_module,
    inputs = [ coremltools.TensorType(shape = x.shape) ],
)
input_name = test_module_mlmodel._spec.description.input[0].name
output_name = test_module_mlmodel._spec.description.output[0].name
mlmodel_output = test_module_mlmodel.predict({input_name: x.detach().numpy()})[output_name]
import numpy as np
np.testing.assert_array_almost_equal(mlmodel_output, pytorch_output, decimal= 3)

System Information

torch: 1.9.1
coremltools: 4.1

apple / coremltools

Got different result using AdaptiveInstanceNorm2d #1357

❓Question

System Information