Pytorch: Add input normalization to model (division layer)

user_2191652 0
Viewed 3481

I want to add the image normalization to an existing pytorch model, so that I don't have to normalize the input image anymore.

Say I have an existing model

model = torch.hub.load('pytorch/vision:v0.6.0', 'mobilenet_v2', pretrained=True)
model.eval()

Now I can add new layers (for example a relu) using torch.nn.Sequential:

new_model = nn.Sequential(
    model,
    nn.ReLU()
)

However I couldn't find a layer to do perform just a division or subtraction as needed for the input normalization here shown in numpy:

import cv2
import numpy as np
img = cv2.imread("my_img.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.astype(np.float32)
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
img = img / 255.0
img = img - mean
img = img / std
img = np.transpose(img, (2, 0, 1))
img = np.expand_dims(img, axis=0)

The goal is that normalization is eventually done on GPU to save time during inference. Also I cannot use torchvision transforms as those operation are not stored inside the model itself. For example, if I want to save the model to disk (in order to convert it to tflite using onnx) the torchvision transform operations will not be saved along with the model. Is there an elegant way of doing this?

(preferably without using a linear layer, which would fix my model input size, which should be flexible as my real model is fully convolutional)

pytorch
2 Answers

Untested code which hopefully you can vet yourself.

import torch.nn as nn

cuda0 = torch.device('cuda:0')

class Normalize(nn.Module):
    def __init__(self, mean, std):
        super(Normlize, self).__init__()
        self.mean = torch.tensor(mean, device=cuda0)
        self.std = torch.tensor(std, device=cuda0)

    def forward(self, input):
        x = input / 255.0
        x = x - self.mean
        x = x / self.std
        return x

In your model you can do

new_model = nn.Sequential(
    Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    model,
    nn.ReLU()
)
fung_82800

The right way of doing this in PyTorch is using dataset transformations. In your specific case, you need torchvision transforms. You can see an example here or here . Copying some part of the code here, for completeness

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
user_13538610
Related