加载预训练的 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 的 zip 文件，其中包含 /models 和一些其他帮助模块。我能够在 Colab 中加载模型，方法是下载 zip，将其展开到我的 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')