Pytorch: Add input normalization to model (division layer)
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(&quot;my_img.jpg&quot;)
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)

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()
)

