How to add preprocess into the model

4

It is possible as shown here:

class MyModel(nn.Module):
    def __init__(self, transform):
        super(MyModel, self).__init__()
        self.conv = nn.Conv2d(3, 1, 3, 1, 1)
        self.transform = transform
        
    def forward(self, x):
        xs = []
        for x_ in x:
            x_ = self.transform(x_)
            xs.append(x_)
        xs = torch.stack(xs)
        x = self.conv(xs)
        return x

transform = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
model = MyModel(transform)
x = torch.randn(1, 3, 24, 24)
output = model(x)

If you were using normalization in your PyTorch model, you have to use the same preprocessing pipeline in your CoreML model.

1 Like