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.

Asked By: Mohny

||

Source

Answer 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.

Answered By: hus

Answer 2

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 By: alexheat

Answer 3

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

import torch

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

Answered By: Yonisen

How to load custom model in pytorch

Question:

Answers: