4. 模型权重文件加载与保存

任务目标

本案例完成YOLOv11权重文件的加载与格式转换闭环：使用PyTorch加载官方预训练权重，检查模型结构与权重键名，并将权重保存为PyTorch原生格式（.pt）与ONNX格式（.onnx）。通过文件生成与再加载验证保存结果正确。

环境依赖

需要 ultralytics（含YOLOv11）、torch、onnx。安装命令：pip install ultralytics onnx。若需导出ONNX，还需安装onnxruntime用于后续验证（可选）。

关键完整代码

import torch
from ultralytics import YOLO

# 1. 加载官方预训练权重（自动下载yolo11n.pt）
model = YOLO("yolo11n.pt")
print("模型类型:", type(model))
print("任务类型:", model.task)

# 2. 获取底层PyTorch模型并检查权重键名
pt_model = model.model
state_dict = pt_model.state_dict()
print("权重键数量:", len(state_dict))
# 打印前5个键名，确认结构
for i, k in enumerate(state_dict.keys()):
    if i < 5:
        print("键名示例:", k)

# 3. 保存为PyTorch格式（覆盖默认保存路径）
save_pt_path = "yolo11n_custom.pt"
torch.save(state_dict, save_pt_path)
print("PyTorch权重已保存至:", save_pt_path)

# 4. 保存为ONNX格式（动态batch，简化输出）
save_onnx_path = "yolo11n_custom.onnx"
model.export(format="onnx", dynamic=True, simplify=True, imgsz=640)
# export默认生成yolo11n.onnx，重命名以区分
import os
if os.path.exists("yolo11n.onnx"):
    os.rename("yolo11n.onnx", save_onnx_path)
print("ONNX权重已保存至:", save_onnx_path)

# 5. 验证：重新加载保存的.pt文件并推理
loaded_model = YOLO(save_pt_path)
print("重新加载成功，任务类型:", loaded_model.task)

# 6. 验证：加载ONNX文件（使用onnxruntime或ultralytics）
onnx_model = YOLO(save_onnx_path)
print("ONNX加载成功，格式:", onnx_model.task)

运行方式

将上述代码保存为 weight_io.py，在Python 3.9+环境中直接运行：

python weight_io.py

首次运行会自动下载 yolo11n.pt（约5.4MB），需保持网络畅通。

验证方法与判定标准

• 文件生成检查：运行结束后，当前目录应存在 yolo11n_custom.pt 与 yolo11n_custom.onnx 两个文件，且大小均大于1MB。
• 加载验证：控制台输出“重新加载成功”与“ONNX加载成功”即表示两种格式均可被YOLO API正常读取。
• 结构一致性：打印的权重键名以 model. 开头（如 model.0.conv.weight），键数量约为300个左右，说明预训练权重完整加载。

常见问题

• ONNX导出失败：若提示缺少onnxsim，运行 pip install onnxsim 后重试；若因网络问题下载超时，可先手动下载 yolo11n.pt 放置于脚本同目录。
• 内存不足：导出ONNX时若显存不足，将 imgsz 改为 320 可降低资源占用。
• 格式差异说明：torch.save 保存的是纯权重字典（不含模型结构），加载时必须配合YOLO类使用；ONNX文件则包含完整计算图，可脱离PyTorch环境部署。