加载预训练的pytorch .pt模型后：ModuleNotFoundError: 找不到'models'模块

Question

我从这个链接下载了一个预训练的yolo模型：https://github.com/WongKinYiu/yolov7/releases

我下载的是yolov7-tiny.pt这个文件。然后我尝试运行代码来加载这个模型，并把它转换成onnx文件：

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)

但是在这里出现了一个错误：

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

错误信息是：

ModuleNotFoundError: No module named 'models'

即使在Colab上也能重现这个问题。我的步骤有什么问题吗？

Answer 1

在你提供的那个仓库里，有一个叫做 Source code 的压缩文件，里面包含了 /models 以及一些其他的辅助模块。我通过下载这个压缩文件，把它解压到我 Google Drive 里一个叫 yolov7 的文件夹中，然后把 yolov7-tiny.pt 文件移动到这个文件夹里，最后运行了以下代码：

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