将自定义训练模型权重转TorchScript时遇AttributeError问题求助

问题场景 #

尝试将自定义训练的YOLOv5模型权重转换为TorchScript格式以集成到Flutter应用中，运行以下脚本时出现错误：

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile
model = torch.load('best.pt',map_location="cpu")
model.eval()
example = torch.rand(1, 3, 640, 640)
traced_script_module = torch.jit.trace(model, example)
optimized_traced_model = optimize_for_mobile(traced_script_module)
optimized_traced_model._save_for_lite_interpreter("best_flutter.pt")

错误信息：

AttributeError: 'dict' object has no attribute 'eval'

错误原因 #

YOLOv5训练生成的best.pt文件并非直接的PyTorch模型实例，而是一个字典结构，里面包含了模型权重、训练超参数、配置信息等内容。直接用torch.load()加载得到的是这个字典，而字典没有.eval()方法，因此调用该方法时触发错误。

解决方法 #

需要使用YOLOv5提供的专用加载函数来正确初始化模型并加载权重，而不是直接解析.pt文件。推荐使用attempt_load函数（YOLOv5内置），修改后的脚本如下：

import torch
from models.experimental import attempt_load
from torch.utils.mobile_optimizer import optimize_for_mobile

# 加载模型，autoshape=False避免自动添加预处理层（适配TorchScript追踪）
model = attempt_load('best.pt', map_location='cpu', autoshape=False)
model.eval()

# 生成示例输入用于TorchScript追踪
example = torch.rand(1, 3, 640, 640)
traced_script_module = torch.jit.trace(model, example)

# 优化模型并保存为移动端可用格式
optimized_traced_model = optimize_for_mobile(traced_script_module)
optimized_traced_model._save_for_lite_interpreter("best_flutter.pt")

说明 #

attempt_load函数会自动解析best.pt中的字典内容，提取模型权重并正确初始化YOLOv5模型实例，返回的对象是标准的PyTorch模型，因此可以正常调用.eval()、.jit.trace()等方法。

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