1. YOLOv5多任务训练框架解析:检测与分割联合优化
这个训练脚本实现了一个基于YOLOv5架构的多任务学习框架,同时支持目标检测和语义分割任务。代码的核心在于通过共享骨干网络(backbone)提取特征,然后分别通过不同的任务头(head)完成检测和分割预测。这种设计能够有效利用计算资源,同时提升模型在多任务场景下的表现。
1.1 核心架构设计
脚本中使用的模型架构定义在models/yolov11_custom_seg_big.yaml配置文件中,从代码逻辑可以看出这是一个支持多任务输出的改进版YOLOv5:
python复制model = Model(model_yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)
关键设计特点包括:
- 共享的骨干网络用于特征提取
- 独立的检测头和分割头用于任务特定预测
- 可选的辅助输出层(auxiliary heads)提升分割性能
- 灵活的权重加载机制支持部分加载(如仅加载backbone)
1.2 多任务损失函数
训练过程中同时计算检测损失和分割损失:
python复制compute_loss = PoseLoss(model) # 检测损失
compute_seg_loss = OhemCELoss(thresh=0.7, ignore_index=-1, aux=False).cuda() # 分割损失
损失权重通过三个系数动态调整:
python复制detgain, seggain, segrm_gain = 0.45, 0.10, 0.45 # 检测、分割比例
这种设计使得模型能够平衡不同任务的学习进度,避免单一任务主导训练过程。
2. 模型初始化与权重加载策略
2.1 灵活的权重加载机制
代码实现了精细化的权重加载控制,支持多种加载模式:
python复制def _load_one_weights(model, ckpt_path, exclude, key_filter=None):
ckpt = torch.load(ckpt_path, map_location=device)
state_dict = _extract_state_dict(ckpt)
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)
if key_filter is not None:
state_dict = {k: v for k, v in state_dict.items() if key_filter(k)}
model.load_state_dict(state_dict, strict=False)
关键特性包括:
- 支持严格模式(strict_load)和宽松模式加载
- 可通过key_filter选择性加载特定层(如仅加载backbone)
- 自动跳过形状不匹配的参数
- 支持多权重文件组合加载
2.2 Backbone专用加载策略
对于只需要加载骨干网络权重的场景,代码提供了专门的过滤函数:
python复制def _make_backbone_key_filter(backbone_len):
def _filter(k):
if not k.startswith('model.'):
return False
parts = k.split('.')
if len(parts) < 3:
return False
try:
idx = int(parts[1])
except ValueError:
return False
return idx < backbone_len
return _filter
这种设计在迁移学习场景中特别有用,可以:
- 使用预训练骨干网络加速收敛
- 保持任务特定头部的随机初始化
- 避免预训练权重对特定任务头部的负面影响
3. 训练流程与优化策略
3.1 混合精度训练配置
脚本使用了PyTorch的AMP(Automatic Mixed Precision)模块实现混合精度训练:
python复制scaler = amp.GradScaler(enabled=cuda)
with amp.autocast(enabled=cuda):
pred = model(imgs)
loss, loss_items = compute_loss(pred[0], targets.to(device))
scaler.scale(loss).backward()
关键配置要点:
- 前向计算在autocast上下文中自动选择合适精度
- GradScaler管理梯度缩放防止下溢出
- 显著减少显存占用同时保持模型精度
3.2 学习率调度与热身
训练采用了组合式的学习率调整策略:
python复制if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # 线性
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
同时包含完善的热身阶段:
python复制if ni <= nw: # warmup阶段
xi = [0, nw]
accumulate = max(1, np.interp(ni, xi, [1, math.floor(nbs / total_batch_size)]).round())
for j, x in enumerate(optimizer.param_groups):
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
3.3 参数分组与优化器配置
代码将模型参数分为三组分别优化:
python复制pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # 偏置参数
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # BN层权重(无衰减)
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # 常规权重(带衰减)
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})
optimizer.add_param_group({'params': pg2})
这种分组策略允许:
- 对BN层使用无衰减优化
- 对常规权重应用L2正则化
- 单独控制偏置参数的学习率
4. 数据加载与增强策略
4.1 多数据集并行加载
脚本支持同时加载检测和分割数据集:
python复制# 检测数据集
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True)
# 分割数据集
seg_trainloader = SegmentationDataset.get_custom_loader(root=segtrain_path, split="train",
base_size=imgsz, batch_size=int(batch_size - 8))
# 另一分割任务数据集
seg_rm_train_dataloder = SegmentationDataset.get_rm_loader(root=seg_rm_train_path,
split="train", base_size=imgsz, batch_size=int(batch_size + 24))
4.2 动态批次组合策略
训练时采用动态批次组合策略平衡不同任务:
python复制seg_dataloder = iter(segpbar)
seg_rm_dataloder = iter(seg_rm_pbar)
for det_batch in pbar:
i, (imgs, targets, paths, _) = det_batch
try:
_, (segimgs, segtargets) = next(seg_dataloder)
except StopIteration:
seg_dataloder = iter(enumerate((seg_trainloader)))
_, (segimgs, segtargets) = next(seg_dataloder)
try:
_, (seg_rm_imgs, seg_rm_targets) = next(seg_rm_dataloder)
except StopIteration:
seg_rm_dataloder = iter(enumerate((seg_rm_train_dataloder)))
_, (seg_rm_imgs, seg_rm_targets) = next(seg_rm_dataloder)
这种设计确保:
- 检测数据持续流动
- 分割数据循环补充
- 不同任务数据比例保持稳定
5. 模型验证与保存策略
5.1 多指标验证体系
验证阶段同时评估检测和分割性能:
python复制# 分割任务验证
mIoU = test.seg_validation(model=ema.ema, valloader=seg_valloader, device=device, n_segcls=2)
mIoU_rm = test.seg_validation(model=ema.ema, valloader=seg_rm_val_dataloder, device=device, n_segcls=14)
# 检测任务验证
results, maps, times = test.test(data_dict, batch_size=batch_size * 2, imgsz=imgsz_test,
model=ema.ema, dataloader=testloader, save_dir=save_dir)
5.2 模型保存与EMA
使用指数移动平均(EMA)提升模型鲁棒性:
python复制ema = ModelEMA(model) if rank in [-1, 0] else None
# 保存检查点
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(model.module if is_parallel(model) else model).half(),
'ema': deepcopy(ema.ema).half(),
'optimizer': optimizer.state_dict()}
torch.save(ckpt, last)
保存策略特点:
- 同时保存原始模型和EMA模型
- 定期保存中间检查点(epoch > 70)
- 自动维护最佳模型副本
- 支持W&B实验跟踪集成
6. 实际训练中的经验技巧
6.1 学习率与批次大小协调
在多任务训练中,不同任务对学习率的敏感度不同。实践中发现:
python复制# 检测与分割损失比例需要根据batchsize调整
# batchsize=13时:0.65, 0.35
# batchsize增大时需降低分割比例或调小学习率
detgain, seggain, segrm_gain = 0.45, 0.10, 0.45
这是因为:
- 分割任务通常需要更稳定的梯度
- 大批次训练时分割损失容易主导优化方向
- 需要通过损失权重维持任务平衡
6.2 梯度累积技巧
脚本实现了动态梯度累积策略:
python复制accumulate = max(round(nbs / total_batch_size), 1) # 累积次数
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # 调整权重衰减
if ni % accumulate == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
关键点:
- 自动计算所需累积步数
- 同步调整权重衰减系数
- 只在累积完成后执行参数更新
- 热身阶段逐步增加累积次数
6.3 分割损失函数选择
代码尝试了多种分割损失函数:
python复制# 常规交叉熵
# compute_seg_loss = SegmentationLosses(aux=False, ignore_index=-1).cuda()
# Focal Loss(城市景观效果不佳)
# compute_seg_loss = SegFocalLoss(ignore_index=-1, gamma=1, reduction="mean").cuda()
# OHEM交叉熵(最终选择)
compute_seg_loss = OhemCELoss(thresh=0.7, ignore_index=-1, aux=False).cuda()
选择依据:
- 常规CE训练稳定但收敛慢
- Focal Loss对类别不平衡敏感但城市景观效果不佳
- OHEM通过困难样本挖掘加速收敛
7. 常见问题与解决方案
7.1 显存不足问题
当遇到显存不足时,可以尝试以下调整:
- 减小batch size(需同步调整梯度累积次数)
- 使用更小的输入尺寸(imgsz)
- 禁用部分数据增强
- 减少模型复杂度(修改yaml配置)
7.2 训练不收敛排查
如果训练损失不下降,建议检查:
- 学习率是否合适(可从1e-3尝试)
- 损失权重比例是否平衡
- 数据加载是否正确(可视化样本检查)
- 模型初始化是否合理(加载预训练权重测试)
7.3 多GPU训练同步问题
使用DDP模式时需注意:
python复制if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
常见问题包括:
- BN层统计不同步(启用SyncBatchNorm)
- 梯度不同步(检查find_unused_parameters)
- 验证集重复评估(确保只在rank=0执行)
7.4 模型部署量化准备
虽然当前脚本未启用量化,但为部署做准备:
python复制# 可选的量化准备
# import torch.ao.quantization as quantizer
建议在模型稳定后:
- 进行训练后量化(PTQ)
- 或实现量化感知训练(QAT)
- 测试量化模型精度下降是否可接受
这个多任务训练框架展示了如何在一个统一模型中高效实现目标检测和语义分割。通过共享特征提取层和任务特定预测头的设计,既节省了计算资源,又保持了各任务的性能。实际应用中可根据具体需求调整模型结构、损失权重和训练策略。
