пример выполнения простого вывода с помощью pytorch-lightning

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

Итак, в настоящее время мой __init__ метод для модели выглядит так:

self._load_config_file(cfg_file)
# just creates the pytorch network
self.create_network()  

self.load_weights(weights_file)

self.cuda(device=0)  # assumes GPU and uses one. This is probably suboptimal
self.eval()  # inference mode

То, что я могу почерпнуть из документации Lightning, я могу сделать то же самое, за исключением того, что cuda()вызов. Так что-то вроде:

self.create_network()

self.load_weights(weights_file)
self.freeze()  # inference mode

Итак, мой первый вопрос: правильно ли это использовать молнию? Как Lightning узнает, нужно ли использовать графический процессор? Я предполагаю, что это нужно где-то указать.

Теперь, для вывода, у меня есть следующая установка:

def infer(frame):
    img = transform(frame)  # apply some transformation to the input
    img = torch.from_numpy(img).float().unsqueeze(0).cuda(device=0)
    with torch.no_grad():
        output = self.__call__(Variable(img)).data.cpu().numpy()
    return output

Это то, что меня смущает. Какие функции мне нужно переопределить, чтобы сделать вывод, совместимый с молнией?

Кроме того, на данный момент ввод поступает в виде массива numpy. Это возможно с помощью модуля Lightning или всегда нужно использовать какой-то загрузчик данных?

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

Я использую 0.7.5 с pytorch 1.4.0 на графическом процессоре с cuda 10.1

Источник

2 ответа

Решение

LightningModule является подклассом torch.nn.Moduleпоэтому один и тот же класс модели будет работать как для вывода, так и для обучения. По этой причине вам, вероятно, следует позвонить вcuda() а также eval() методы за пределами __init__.

Поскольку это просто nn.ModuleПод капотом, как только вы загрузили свои веса, вам не нужно переопределять какие-либо методы для выполнения вывода, просто вызовите экземпляр модели. Вот пример игрушки, которую вы можете использовать:

import torchvision.models as models
from pytorch_lightning.core import LightningModule

class MyModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.resnet = models.resnet18(pretrained=True, progress=False)
    
    def forward(self, x):
        return self.resnet(x)

model = MyModel().eval().cuda(device=0)

А затем, чтобы фактически выполнить вывод, вам не нужен метод, просто выполните что-то вроде:

for frame in video:
    img = transform(frame)
    img = torch.from_numpy(img).float().unsqueeze(0).cuda(0)
    output = model(img).data.cpu().numpy()
    # Do something with the output

Основное преимущество PyTorchLighting заключается в том, что вы также можете использовать тот же класс для обучения, реализовав training_step(), configure_optimizers() а также train_dataloader()в этом классе. Вы можете найти простой пример этого в документации PyTorchLightning.

Источник

Несмотря на то, что приведенного выше ответа достаточно, если принять к сведению следующую строку

      img = torch.from_numpy(img).float().unsqueeze(0).cuda(0)

Нужно поместить как модель, так и изображение в правильный графический процессор. На машине логического вывода с несколькими графическими процессорами это становится проблемой.

Чтобы решить эту проблему,.predictтакже недавно был выпущен, см. больше на https://pytorch-lightning.readthedocs.io/en/stable/deploy/production_basic.html .

Источник
Другие вопросы по тегам