Home > Software engineering >  How to use trained pytorch model for prediction
How to use trained pytorch model for prediction

Time:08-18

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 am new to pytorch and deep learning so i am having trouble with available solutions so i need a detailed guide. Thanks

CodePudding user response:

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.

CodePudding user response:

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
  • Related