3

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.

Mohny
  • 49
  • 1
  • 2
  • 7

3 Answers3

10

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

alexheat
  • 479
  • 5
  • 9
1

In order to load your model's weights, you should first import your model script. I guess it is located in /weights/last.py. Afterwards, you can load your model's weights.

Example code might be as below:

import os  
import torch 
from weights.last import Model # I assume you named your model as Model, change it accordingly

model = Model()  # Then in here instantiate your model
model.load_state_dict(torch.load(ospath.join(os.getcwd()+'/weights/last.pt')))  # Then load your model's weights.

 

# 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)

In this solution, do not forget that you should run your program from your current working directory, if you run it from the weights folder you might receive errors.

hus
  • 56
  • 5
  • Thanks! I used yolov5 for training. Where can I find my model? I cant see in weights folder any other data. – Mohny Nov 30 '21 at 12:34
  • I do not know Yolov5 but if you got the source code from https://github.com/ultralytics/yolov5, I think it is located as models/yolov5.py – hus Nov 30 '21 at 12:48
  • thanks @hus. But I really still couldn't get it worked. I can't see py file in model folder. Do you have any idea where or how I can find it? I saw the folder https://github.com/ultralytics/yolov5/tree/master/models here but it also didn't work. – Mohny Nov 30 '21 at 23:17
-1

If you wanna load your local saved model you can try this

import torch

model = torch.hub.load('.', 'custom', 'yourmodel.pt', source='local')
Yonisen
  • 1
  • 1