CODEWITHSUNDEEP
pytorch
mcExchange
mcExchange

Reputation: 6496

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

Upvotes: 3

Views: 7182

Answers (1)

Fung
Fung

Reputation: 7760

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

Upvotes: 7

Related Questions