如何解决PyTorch加载自定义YOLOv5n模型时torch.load触发的ModuleNotFoundError

错误原因 #

你在执行torch.load()时抛出ModuleNotFoundError的核心原因是：PyTorch的torch.load()反序列化权重时需要依赖训练时的模型类定义，你当前运行环境没有引入YOLOv5对应的结构模块，无法识别权重内嵌的类路径。

修复方案 #

方案1：使用YOLOv5官方接口加载（推荐） #

该方法会自动匹配对应的模型结构定义，无需手动引入结构文件，修改后代码如下：

import os
import torch

# 加载自定义训练的YOLOv5权重
model = torch.hub.load('ultralytics/yolov5', 'custom', path=os.getcwd()+'/weights/last.pt', force_reload=False)

# 推理逻辑保持不变
imgs = ['https://example.com/img.jpg']   
# 执行推理
results = model(imgs)
# 结果处理
results.print()
results.save()  # 或调用.show()展示
results.xyxy[0]  # 获取tensor格式的预测结果
results.pandas().xyxy[0]  # 获取pandas格式的预测结果

方案2：手动引入本地结构文件后加载 #

如果你要沿用torch.load()的写法，需要先将本地存放的模型结构文件加入运行环境，修改后代码如下：

import os
import torch
import sys
# 将结构文件所在目录加入Python运行路径
sys.path.append(os.getcwd()+'/weights')
# 引入结构文件中的模型类，注意替换为你last.py中实际的类名
from last import YOLOv5Model

# 实例化模型后加载权重参数
model = YOLOv5Model()
checkpoint = torch.load(os.getcwd()+'/weights/last.pt')
# 若权重仅保存参数，使用load_state_dict加载
model.load_state_dict(checkpoint)
# 切换为推理模式
model.eval()

# 推理逻辑保持不变
imgs = ['https://example.com/img.jpg']   
results = model(imgs)
results.print()
results.save()
results.xyxy[0]
results.pandas().xyxy[0]

注意：如果你的last.pt是训练时保存的完整模型（同时包含结构和参数），则无需调用load_state_dict，但仍需提前引入last.py中的模型类，否则反序列化时仍会报模块不存在错误。

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