Pytorch: добавить в модель нормализацию ввода (слой разделения) - PullRequest
Новая
       3

Pytorch: добавить в модель нормализацию ввода (слой разделения)

1 голос
/ 19 июня 2020

Я хочу добавить нормализацию изображения к существующей модели pytorch, чтобы мне больше не приходилось нормализовать входное изображение.

Допустим, у меня есть существующая модель 

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

Теперь я могу добавлять новые слои (например, relu) с помощью torch.nn.Sequential: 

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, поскольку эти операции не хранятся внутри самой модели. Например, если я хочу сохранить модель на диск (чтобы преобразовать ее в tflite с помощью onnx), операции преобразования torchvision не будут сохранены вместе с моделью. Есть ли элегантный способ сделать это?

(желательно без использования линейного слоя, который исправит размер ввода моей модели, который должен быть гибким, поскольку моя реальная модель полностью сверточная)

Ответы [ 2 ]

1 голос
/ 24 июня 2020

Непроверенный код, который, надеюсь, вы сможете проверить самостоятельно. 

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()
)
0 голосов
/ 20 июня 2020

Правильный способ сделать это в PyTorch - использовать преобразования набора данных. В вашем конкретном случае c вам нужно преобразований torchvision . Вы можете увидеть пример здесь или здесь . Скопируйте сюда часть кода, для полноты 

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