如何使用经过训练的Pytorch模型进行预测

y0u0uwnf  于 12个月前  发布在  其他
关注(0)|答案(2)|浏览(146)

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

brtdzjyr

brtdzjyr1#

要使用预训练模型,您应该将状态加载到架构的新示例上,如docs/tutorials中所述:
这里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示例中。
  • 用数据做出预测。
  • 计算结果的指标。

更多关于DatasetDataLoaderhere

1zmg4dgp

1zmg4dgp2#

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

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

字符串

相关问题