如何在PyTorch中修改YOLOv5/YOLOv7训练模型的类别名称并保存

当然可以修改类别名称后重新导出成PyTorch的.pt模型！你已经走完了最核心的一步——加载模型并更新names属性，接下来只需要用PyTorch内置的保存方法就能搞定了。

直接保存修改后的模型 #

如果你的模型加载后直接是模型实例（像你代码里那样能直接设置model.names），直接用torch.save()就行：

import torch

# 加载原模型，指定设备（cpu/cuda）
map_location = torch.device('cpu')  # 按需切换成'cuda'
model = torch.load('your_original_model.pt', map_location=map_location)

# 更新类别名称列表
model.names = ['face', 'head', 'helmet']

# 保存修改后的模型为best.pt
torch.save(model, 'best.pt')

特殊情况：加载后是字典结构 #

有些YOLO版本会把模型、训练状态、优化器等信息打包成字典保存，这时候需要先取出模型对象修改，再放回字典保存：

import torch

map_location = torch.device('cpu')
model_dict = torch.load('your_original_model.pt', map_location=map_location)

# 修改字典里的模型类别名称
model_dict['model'].names = ['face', 'head', 'helmet']

# 保存完整的字典结构
torch.save(model_dict, 'best.pt')

验证修改是否生效 #

保存完成后，可以重新加载模型检查一下，确保类别名称已经更新：

saved_model = torch.load('best.pt', map_location=map_location)
print(saved_model.names)  # 应该输出['face', 'head', 'helmet']

要提醒你的是：这个操作只会修改类别名称的映射关系，完全不会触动模型的权重参数，所以检测性能不会受到任何影响，只是推理时输出的类别标签会变成你设置的新名称~

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