Below is the source code, I use to load a .pth
file and do a multi-class image classification prediction.
model = Classifier() # The Model Class.
model.load_state_dict(torch.load('<PTH-FILE-HERE>.pth'))
model = model.to(device)
model.eval()
# prediction function to test images
def predict(img_path):
image = Image.open(img_path)
resize = transforms.Compose(
[ transforms.Resize((256,256)), transforms.ToTensor()])
image = resize(image)
image = image.to(device)
y_result = model(image.unsqueeze(0))
result_idx = y_result.argmax(dim=1)
print(result_idx)
I converted the .pth
file to an ONNX file using torch.onnx.export
.
Now, How can I write a prediction script similar to above one by using the ONNX file alone and not using the .pth
file.?
Is it possible to do so?
You can use ONNX Runtime.
You can follow the tutorial for detailed explanation.
Usually, the purpose of using onnx is to load the model in a different framework and run inference there e.g. PyTorch -> ONNX -> TensorRT.