Search code examples
pythonpytorchcomputer-visionyolov5

How to use pt file

I'm trying to make a currency recognition model and I did so using a dataset on kaggle and colab using yolov5 and I exactly carried out the steps explained on yolov5 github. At the end, I downloaded a .pt file which has the weights of the model and now I want to use it in python file to detect and recognize currency . How to do this?

I am a beginner in computer vision and I am totally confused about what to do. I am searching over and over but I don't reach anything.

import torch

# Model
model=torch.load('E:\_best.pt')

# Images
imgs=['E:\Study\currency.jpg']

# Inference
results = model(imgs)

# Results
results.print()
results.save()  # or .show()
results.show()
results.xyxy[0]  # img1 predictions (tensor)
results.pandas().xyxy[0]
Solution

  • If you want to read trained parameters from .pt file and load it into your model, you could do the following.

    file = "model.pt"
model = your_model()
model.load_state_dict(torch.load(file))
# this will automatically load the file and load the parameters into the model.

    before calling load_state_dict(), be sure that the .pt file contains only model parameters, otherwise, error occurs. This can be checked by print(torch.load(file)).