Search code examples
pythonpytorchyoloyolov5

How to load custom model in pytorch

I'm trying to load my pretrained model (yolov5n) and test it with the following code in PyTorch:

import os  
import torch 
model = torch.load(os.getcwd()+'/weights/last.pt')
 

# Images
imgs = ['https://example.com/img.jpg']   
# Inference
results = model(imgs)

# Results
results.print()
results.save()  # or .show()

results.xyxy[0]  # img1 predictions (tensor)
results.pandas().xyxy[0]  # img1 predictions (pandas)

and I'm getting the following error:

ModuleNotFoundError Traceback (most recent call last) in 3 import torch 4 ----> 5 model = torch.load(os.getcwd()+'/weights/last.pt')

My model is located in the folder /weights/last.py, I'm not sure what I'm doing false. Could you please tell me, what it's missing in my code.

Solution

  • You should be able to find the weights in this directory: yolov5/runs/train/exp/weights/last.pt

    Then you load the weights with a line like this:

    model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/runs/train/exp/weights/last.pt', force_reload=True)

    I have an example of a notebook that loads custom models and from that directory after training the model here https://github.com/pylabel-project/samples/blob/main/pylabeler.ipynb