如何使用经过训练的 pytorch model 进行预测

Question

I have a pretrained pytorch model which is saved in.pth format. How can i use it for prediction on new dataset in a separate python file.

I need a detailed guide.

Answer 1

To use a pretrained model you should load the state on a new instance of the architecture as explained in the docs/tutorials :

Here models is imported beforehand:

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.

If you are using a custom architecture you only need to change the first line.

model = MyCustomModel()

After enabling the eval mode, you can proceed as follows:

• Load your data into a Dataset instance and then in a DataLoader .
• Make your predictions with the data.
• Calculate metrics on the results.

More about Dataset and DataLoader here .

Answer 2

Well for prediction theres something called forward pass

import torch
from torch_model import Model # Made up package

device = torch.device('cpu' if torch.cuda.is_available() else 'gpu')

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