Revision history [back]

To obtain predictions from a PyTorch model, perform the following steps:

1. Load the PyTorch model using torch.load() or torch.load_state_dict() functions.
2. Set the model to eval() mode to disable dropout and batch normalization layers. This ensures that the model behaves consistently during inference.
3. Convert the data to the format expected by the model. For example, if the model expects images with 3 channels of 224x224 pixels, convert input images to this format.
4. Pass the input data to the model to obtain the output tensor, using the model.forward() function.
5. Depending on the output format, extract the required information from the output tensor. For example, if the model outputs probabilities for a set of classes, apply softmax to the output tensor to obtain the probabilities and select the class with the highest probability as the predicted class.

Here is some sample code to illustrate these steps:

import torch
from torchvision.transforms import ToTensor
from PIL import Image

# Step 1: Load the PyTorch model
model = torch.load('model.pth')

# Step 2: Set the model to eval mode
model.eval()

# Step 3: Convert the input data to the expected format
img = Image.open('input.png')
img = ToTensor()(img)
img = img.unsqueeze(0) # Add batch dimension

# Step 4: Pass the input data to the model to obtain the output tensor
output = model.forward(img)

# Step 5: Extract the required information from the output tensor
probs = torch.nn.functional.softmax(output, dim=1)
pred_class = torch.argmax(probs)

print(pred_class)