torch.load加载doclayout_yolo.pt模型权重失败的解决方法

问题场景 #

已将模型文件doclayout_yolo.pt与代码文件weights-script.py放在同一目录，且通过以下命令安装依赖：

python -m pip install doclayout-yolo 
python -m pip install huggingface_hub  

尝试用自定义的ExampleModel加载权重时，先后出现两类错误：

1. 用weights_only=False加载时，提示权重键不匹配，存在多余的训练状态参数（如epoch、optimizer）且缺少MLP模型的权重键。
2. 改用weights_only=True加载时，提示PyTorch默认不允许加载doclayout_yolo.nn.tasks.YOLOv10DetectionModel类。

错误原因 #

• 自定义的ExampleModel是简单的MLP结构，和doclayout_yolo.pt中存储的YOLOv10检测模型结构完全不匹配，权重键自然无法对应。
• doclayout_yolo.pt是完整的训练 checkpoint，包含模型权重、训练状态、优化器参数等，并非适配自定义MLP的权重文件。
• 用weights_only=True时，PyTorch默认限制加载非内置的模型类，需要将目标模型类加入安全白名单。

修复方案 #

方案1：直接使用库内置模型类加载（推荐） #

替换原有代码，用doclayout-yolo提供的模型类加载权重：

import torch
from doclayout_yolo.nn.tasks import YOLOv10DetectionModel

# 方式1：用from_pretrained一键加载
model = YOLOv10DetectionModel.from_pretrained('doclayout_yolo.pt')
model.eval()  # 切换为评估模式

# 方式2：手动加载checkpoint（如果方式1报错）
# checkpoint = torch.load('doclayout_yolo.pt', weights_only=False, map_location='cpu')
# model = YOLOv10DetectionModel(checkpoint['model'].args)
# model.load_state_dict(checkpoint['model'].state_dict())
# model.eval()

print(model)

方案2：使用weights_only=True加载（需添加安全白名单） #

如果必须用weights_only=True，先将模型类加入PyTorch的安全全局列表：

import torch
from doclayout_yolo.nn.tasks import YOLOv10DetectionModel

# 将YOLOv10DetectionModel加入安全白名单
torch.serialization.add_safe_globals([YOLOv10DetectionModel])

# 加载checkpoint
checkpoint = torch.load('doclayout_yolo.pt', weights_only=True, map_location='cpu')
# 初始化对应结构的模型
model = YOLOv10DetectionModel(checkpoint['model'].args)
# 加载模型权重
model.load_state_dict(checkpoint['model'].state_dict())
model.eval()

print(model)

关键说明 #

永远不要用自定义的模型结构去加载专为其他模型训练的权重文件——模型结构和权重键必须严格对应。doclayout_yolo.pt是为YOLOv10检测模型设计的，必须用doclayout-yolo库中的YOLOv10DetectionModel类加载。

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