使用 pytorch-lightning 进行简单预测的示例

Posted 2023-03-27

【中文标题】使用 pytorch-lightning 进行简单预测的示例

【英文标题】：example of doing simple prediction with pytorch-lightning

【发布时间】：2020-08-17 09:14:55

【问题描述】：

我有一个现有模型，我在其中加载一些预训练的权重，然后在 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()  # prediction mode

我可以从闪电文档中收集到的信息，我几乎可以做同样的事情，除了不拨打 cuda() 电话。所以像：

self.create_network()

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

那么，我的第一个问题是这是否是使用闪电的正确方法？闪电如何知道它是否需要使用 GPU？我猜这需要在某个地方指定。

现在，对于预测，我有以下设置：

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 数组。这是闪电模块可以实现的，还是总是需要使用某种数据加载器？

在某些时候，我想扩展这个模型实现来进行训练，所以想确保我做对了，但是虽然大多数例子都集中在训练模型上，一个简单的例子，只是在生产时做预测单个图像/数据点可能有用。

我在带有 cuda 10.1 的 GPU 上使用 0.7.5 和 pytorch 1.4.0

【参考方案1】：

LightningModule 是torch.nn.Module 的子类，因此同一个模型类将适用于推理和训练。因此，您可能应该在 __init__ 之外调用 cuda() 和 eval() 方法。

因为它只是一个 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 docs 中找到一个简单的示例。