Pytorch 使用Pytork lightning进行简单推理的示例

我有一个现有的模型，我加载一些预先训练好的权重，然后在pytorch中进行推理（一次一张图像）。我试图将其转换为pytorch lightning模块，但对一些事情感到困惑

因此，目前，我的模型的

\uuuuu init\uuuu

方法如下所示：

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的正确方法？lightning如何知道是否需要使用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

这一点让我感到困惑。我需要重写哪些函数才能做出与lightning兼容的推断

此外，目前，输入是一个numpy数组。这是可以从lightning模块中实现的，还是必须使用某种数据加载器

在某种程度上，我希望扩展此模型实现以进行培训，因此我希望确保我做得正确，但尽管大多数示例都集中在培训模型上，但一个简单的示例（仅在生产时对单个图像/数据点进行推断）可能会很有用

---

我正在cuda 10.1的GPU上使用0.7.5和pytorch 1.4.0，LightningModule是torch.nn.Module的子类，因此相同的模型类将用于推理和训练。出于这个原因，您可能应该在

\uu init\uu

之外调用

cuda（）

和

eval（）

方法

因为它只是引擎盖下的一个

nn.Module

，一旦加载了权重，就不需要重写任何方法来执行推理，只需调用模型实例即可。以下是您可以使用的玩具示例：

导入torchvision.models作为模型
从pytorch_lightning.core导入lightning模块
类别MyModel（LightningModule）：
定义初始化（自）：
super（）。\uuuu init\uuuuu（）
self.resnet=models.resnet18（pretrained=True，progress=False）
def前进（自身，x）：
返回self.resnet（x）
model=MyModel（）.eval（）.cuda（设备=0）

然后，要实际运行推断，您不需要方法，只需执行以下操作：

self.create_network()

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

视频中帧的

：
img=变换（帧）
img=torch.from_numpy（img）.float（）.unsqueze（0）.cuda（0）
输出=模型（img）.data.cpu（）.numpy（）
#对输出做些什么

PyTorchLighting的主要好处是，通过在该类上实现

training\u step（）

、

configure\u optimizers（）

和

train\u dataloader（）

，您还可以使用同一类进行培训。您可以在中找到一个简单的例子