after loading a pretrained pytorch .pt model file: ModuleNotFoundError: No module named 'models'

Question

I downloaded one of the pretrained yolo models from the link: https://github.com/WongKinYiu/yolov7/releases

In this case, yolov7-tiny.pt is downloaded. Then tried to run the code to load the model and convert it to onnx file:

import torch
import onnx

model = torch.load('./yolo_custom/yolov7-tiny.pt')
input_shape = (1, 3, 640, 640)
torch.onnx.export(model, torch.randn(input_shape), 'yolov7-tiny.onnx', opset_version=11)

An error occurs on

model = torch.load('./yolo_custom/yolov7-tiny.pt')

and the error message is:

ModuleNotFoundError: No module named 'models'

The issue is reproducible even on Colab. Is there anything wrong on the steps?

Answer 1

In the repository you linked, there is a zip file called Source code which contains /models and some other helper modules. I was able to load the model in Colab by downloading the zip, expanding it to a directory in my Google Drive called yolov7, moving yolov7-tiny.pt to this directory and then running the following:

from google.colab import drive
drive.mount('/content/drive')

import torch
import onnx
import sys

sys.path.append('/content/drive/My Drive/yolov7')
model = torch.load('/content/drive/My Drive/yolov7/yolov7-tiny.pt')