How to load custom model in pytorch

Question

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.

alexheat · Accepted Answer · 2021-12-01 04:45:52Z

11

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

answered Dec 1, 2021 at 4:45

alexheat

6046 silver badges11 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

hus · Accepted Answer · 2021-11-30 12:15:04Z

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.

Thanks! I used yolov5 for training. Where can I find my model? I cant see in weights folder any other data.
I do not know Yolov5 but if you got the source code from github.com/ultralytics/yolov5, I think it is located as models/yolov5.py
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 github.com/ultralytics/yolov5/tree/master/models here but it also didn't work.

Yonisen · Accepted Answer · 2022-11-16 13:58:09Z

-2

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

import torch

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

edited Nov 16, 2022 at 13:58

answered Nov 16, 2022 at 13:55

Yonisen

11 bronze badge

Collectives™ on Stack Overflow

How to load custom model in pytorch

3 Answers 3

Comments

3 Comments

Comments

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

3 Comments

Comments

Related