如何在Python中通过PyTorch加载带自定义权重的yolov5l模型

YOLOv5的.pt权重文件在训练保存时，默认序列化了模型类的路径依赖，训练时是在官方YOLOv5项目目录下执行的，权重内记录的模型类归属为models.yolo.Model，你自己的项目没有对应路径的models模块，直接调用torch.load()就会触发模块不存在的报错。

方法1：依赖YOLOv5项目结构加载 #

这个方法适合需要保留PyTorch原生模型可训练/可调整属性的场景：

1. 把你训练时用的YOLOv5项目中的models、utils两个目录完整复制到你自己项目的根目录下
2. 用如下代码加载权重：

import torch
from models.yolo import Model

# 填写和你训练时完全一致的参数
nc = 你的训练类别数
yolov5l_cfg_path = "models/hub/yolov5l.yaml"
weights_path = "你的权重文件路径.pt"

# 初始化模型结构
model = Model(yolov5l_cfg_path, nc=nc)
# 加载权重参数
ckpt = torch.load(weights_path, map_location="cpu")
model.load_state_dict(ckpt["model"].state_dict())
# 切换为推理模式
model.eval()

方法2：导出为通用推理格式（更适合部署场景） #

如果你的场景只需要做推理不需要调整模型结构，可以在训练用的YOLOv5项目里提前把.pt权重导出为ONNX或者TorchScript格式，导出后的文件不需要依赖YOLOv5的项目结构就能直接加载：

1. 在原训练环境的YOLOv5目录下执行导出命令：
   python export.py --weights 你的权重文件路径.pt --include onnx torchscript
2. 你自己的项目里可以直接用onnxruntime、OpenCV dnn或者torch.jit.load加载导出后的文件，不需要复制models、utils目录，也不会出现模块找不到的报错。

• 加载权重时用到的配置文件（如yolov5l.yaml）必须和训练时使用的版本、参数完全一致，否则会出现结构不匹配的报错
• 有GPU推理需求可以把torch.load的map_location参数改为cuda
• 推理前必须调用model.eval()关闭训练专属的dropout、batch norm更新逻辑，否则推理结果会出现异常偏差

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