Pytorch：向模型添加输入归一化（划分层）

Question

我想将图像标准化添加到现有的 pytorch 模型中，这样我就不必再标准化输入图像了。

假设我有一个现有模型

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

现在我可以使用 torch.nn.Sequential 添加新层（例如 relu）：

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

然而，我找不到一个层来根据 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)

目标是归一化最终在 GPU 上完成，以节省推理过程中的时间。另外，我无法使用 torchvision 变换，因为这些操作不存储在模型本身内部。例如，如果我想将模型保存到磁盘（以便使用 onnx 将其转换为 tflite），torchvision 变换操作将不会与模型一起保存。有没有一种优雅的方法来做到这一点？

（最好不使用线性层，这将固定我的模型输入大小，这应该是灵活的，因为我的真实模型是完全卷积的）

Answer 1

未经测试的代码，希望您可以自行审查。

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

在你的模型中你可以做

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