如何使用训练好的pytorch模型进行预测

Question

我有一个预训练的 pytorch 模型，以 .pth 格式保存。我如何使用它来预测单独的 python 文件中的新数据集。

我需要一份详细的指南。

Answer 1

要使用预训练模型，您应该将状态加载到架构的新实例上，如文档/教程中所述：

这里

models

是预先导入的：

model = models.vgg16()
model.load_state_dict(torch.load('model_weights.pth')) # This line uses .load() to read a .pth file and load the network weights on to the architecture.
model.eval() # enabling the eval mode to test with new samples.

如果您使用自定义架构，则只需更改第一行。

model = MyCustomModel()

启用

eval

模式后，您可以进行以下操作：

• 将数据加载到

  Dataset

  实例中，然后加载到

  DataLoader

  中。
• 利用数据做出预测。
• 计算结果指标。

更多关于

Dataset

和

DataLoader

这里。

Answer 2

对于预测来说，有一种叫做前向传播的东西

import torch
from torch_model import Model # Made up package

# select gpu when available, else work with cpu resources
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = Model()
model.load_state_dict(torch.load('weights.pt'))

model = model.to(device) # Set model to gpu
model.eval();

inputs = torch.random.randn(1, 3, 224, 224) # Dtype is fp32
inputs = inputs.to(device) # You can move your input to gpu, torch defaults to cpu

# Run forward pass
with torch.no_grad():
  pred = model(inputs)

# Do something with pred
pred = pred.detach().cpu().numpy() # remove from computational graph to cpu and as numpy