加载YOLOv7预训练.pt模型时出现ModuleNotFoundError: No module named 'models'

问题根源 #

直接用torch.load()加载YOLOv7预训练模型会报错，是因为YOLOv7的.pt模型文件仅保存了模型参数，PyTorch需要对应模型结构的定义（即models模块）才能重构完整模型，而你当前运行环境中没有包含YOLOv7官方代码里的模型结构定义文件。

正确操作步骤 #

1. 获取YOLOv7官方代码
   先把YOLOv7的官方代码仓库内容放到你的运行环境中（本地或Colab），进入代码所在目录。

2. 放置预训练模型
   将下载好的yolov7-tiny.pt放到上述代码目录下（或在代码中指定正确的文件路径）。

3. 使用官方加载函数重构模型并导出ONNX
   替换原代码为以下内容，利用官方提供的attempt_load函数自动处理模型结构加载：

   import torch
   import onnx
   from models.experimental import attempt_load
   
   # 加载预训练模型，自动导入所需模型结构
   model = attempt_load('yolov7-tiny.pt', map_location=torch.device('cpu'))
   model.eval()  # 必须切换到评估模式，否则导出ONNX会有问题
   
   # 构造输入张量并导出ONNX
   input_shape = (1, 3, 640, 640)
   dummy_input = torch.randn(input_shape)
   torch.onnx.export(
       model,
       dummy_input,
       'yolov7-tiny.onnx',
       opset_version=11,
       do_constant_folding=True,
       input_names=['input'],
       output_names=['output']
   )
   

补充说明 #

• 原代码错误在于未加载YOLOv7的模型结构定义，直接加载参数文件无法生成可运行的模型实例。
• 在Colab中操作时，需先确保已获取YOLOv7官方代码并进入对应目录，再执行上述代码。

内容的提问来源于stack exchange，提问作者weiming