YOLO新增类别后能否基于已训练权重续训及报错解决

核心问题回应 #

• 仅用新类别图像续训不可行：YOLO的输出层是针对所有训练类别优化的，只训新类会彻底丢失旧类的识别能力，模型对旧类的检测/分类精度会降到随机水平。
• 可以基于旧模型加新类训练（无需从头开始）：但需要调整模型输出层适配新类别数，同时必须使用包含旧类+新类的完整数据集（或旧类少量样本）进行微调，避免遗忘。

报错原因 #

你遇到的RuntimeError确实是因为新类别数改变了模型输出层的参数形状——原模型输出层是为旧类别数设计的，直接加载到新类别数的模型中，参数维度不匹配，导致加载失败。

具体解决方案 #

方案1：手动适配权重并微调 #

1. 修改配置文件：将cfg/training/yolov7_custom.yaml中的nc参数改为新的类别总数（比如你当前的21类），同时确保data.yaml里的类别数和名称对应正确。
2. 调整权重文件：编写脚本过滤原权重中不匹配的输出层参数，只保留特征提取层的权重，重新初始化输出层：

   import torch
   from models.yolo import Model
   
   # 加载原训练权重
   old_checkpoint = torch.load('runs/train/products/weights/last.pt')
   old_weights = old_checkpoint['model'].state_dict()
   
   # 初始化新模型（基于修改后的配置文件）
   new_model = Model('cfg/training/yolov7_custom.yaml')
   new_state_dict = new_model.state_dict()
   
   # 过滤出形状匹配的参数（跳过输出层）
   filtered_weights = {}
   for k, v in old_weights.items():
       if k in new_state_dict and v.shape == new_state_dict[k].shape:
           filtered_weights[k] = v
   
   # 更新新模型权重，保存为适配后的权重文件
   new_state_dict.update(filtered_weights)
   new_model.load_state_dict(new_state_dict)
   torch.save({'model': new_model.state_dict()}, 'runs/train/products/weights/adapted_last.pt')
   

3. 微调训练：使用适配后的权重，搭配全类别数据集训练，注意用微调专用的超参数：

   !python train.py --workers 8 --batch-size 16 --data products/data.yaml --img 640 640 --cfg cfg/training/yolov7_custom.yaml --epochs 20 --weights 'runs/train/products/weights/adapted_last.pt' --name yolov7_custom_newClass --hyp data/hyp.finetune.yaml --device 0
   

方案2：官方迁移学习流程（更便捷） #

YOLOv7支持自动适配新类别数，通过冻结特征层先训练输出层，再全微调：

1. 修改配置：确保cfg和data.yaml的nc参数为新类别总数。
2. 冻结特征层训练输出层：先冻结大部分特征提取层，只训练输出层适配新类别：

   !python train.py --workers 8 --batch-size 16 --data products/data.yaml --img 640 640 --cfg cfg/training/yolov7_custom.yaml --epochs 5 --weights 'runs/train/products/weights/last.pt' --name yolov7_custom_newClass --hyp data/hyp.finetune.yaml --device 0 --freeze 100
   

   （--freeze 100表示冻结前100层，可根据YOLOv7结构调整，比如官方推荐冻结90层左右）
3. 全层微调：解冻所有层，用较低学习率微调整个模型：

   !python train.py --workers 8 --batch-size 16 --data products/data.yaml --img 640 640 --cfg cfg/training/yolov7_custom.yaml --epochs 20 --weights 'runs/train/yolov7_custom_newClass/weights/last.pt' --name yolov7_custom_full_finetune --hyp data/hyp.finetune.yaml --device 0
   

关键注意事项 #

• 必须使用包含旧类+新类的数据集，否则模型会遗忘旧类的识别能力。
• 微调时使用hyp.finetune.yaml而非hyp.scratch.custom.yaml，这类超参数的学习率更低，避免破坏已训练好的特征提取层。
• 输出层参数无法直接复用，必须重新初始化或通过训练适配新类别。

内容的提问来源于stack exchange，提问作者Afique Amin Zian