王科特1. 项目概述：基于YOLOv5的多任务目标检测与语义分割训练框架
这个项目是一个基于PyTorch实现的计算机视觉训练框架，主要针对多任务学习场景，同时支持目标检测和语义分割任务。核心代码train_608_736.py实现了一个高度可配置的训练流程，能够灵活适应不同硬件环境和数据集需求。
作为一名长期从事计算机视觉开发的工程师，我在实际项目中经常遇到需要同时处理检测和分割任务的场景。这套框架经过多次迭代优化，在保持YOLOv5高效特性的基础上，增加了对语义分割的支持，特别适合道路场景、工业质检等需要多任务输出的应用。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 核心架构解析
2.1 模型结构设计
框架采用的主干网络是基于YOLOv5改进的多任务架构，主要特点包括：
- 共享特征提取：底层特征提取器同时服务于检测和分割任务，减少计算冗余
- 任务特定头：检测头保持YOLOv5原有设计，分割头采用类似DeepLab的结构
- 灵活配置：通过yaml文件定义网络结构，支持不同复杂度的模型变体
模型初始化时，会根据配置文件动态构建网络：
python复制model = Model(model_yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)
2.2 多任务损失函数
框架实现了多种损失函数的组合：
- 检测损失：包含分类损失、定位损失和置信度损失
python复制compute_loss = PoseLoss(model) # 初始化检测criteria
- 分割损失：支持多种选择，包括：
- 标准交叉熵损失
- Focal Loss
- OHEM (Online Hard Example Mining) Loss
python复制compute_seg_loss = OhemCELoss(thresh=0.7, ignore_index=-1, aux=False).cuda()
- 损失权重平衡：通过detgain、seggain等参数控制各任务损失的比例
python复制detgain, seggain, segrm_gain = 0.45, 0.10, 0.45 # 检测,分割比例
3. 训练流程详解
3.1 数据准备与加载
框架支持多种数据加载方式：
- 检测数据加载：继承YOLOv5的数据增强策略
python复制dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect)
- 分割数据加载：专门的分割数据处理器
python复制seg_trainloader = SegmentationDataset.get_custom_loader(root=segtrain_path,
split="train", mode="train",
base_size=imgsz,
batch_size=int(batch_size - 8),
workers=opt.workers, pin=True)
- 数据增强：包括Mosaic、MixUp等YOLOv5特色增强，以及针对分割任务的特定增强
3.2 训练优化策略
训练过程采用了一系列优化技术：
- 学习率调度：支持线性warmup和余弦退火
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']
- 混合精度训练：使用AMP加速训练
python复制scaler = amp.GradScaler(enabled=cuda)
with amp.autocast(enabled=cuda):
pred = model(imgs)
loss, loss_items = compute_loss(pred[0], targets.to(device))
- 梯度累积：支持大batch训练
python复制accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing
- 模型EMA：指数移动平均提升模型稳定性
python复制ema = ModelEMA(model) if rank in [-1, 0] else None
4. 关键实现细节
4.1 多尺度训练
框架支持多尺度训练，增强模型鲁棒性：
python复制if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs
sf = sz / max(imgs.shape[2:])
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
4.2 分布式训练支持
框架完整支持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()))
4.3 模型验证与评估
训练过程中定期进行验证，评估指标包括：
- 检测任务的mAP
- 分割任务的mIoU
python复制if epoch % 5 == 0 or (epochs - epoch) < 60:
mIoU = test.seg_validation(model=ema.ema, valloader=seg_valloader,
device=device, n_segcls=2, half_precision=True)
5. 实战经验与调优技巧
5.1 超参数设置建议
根据实际项目经验，推荐以下超参数配置：
- 学习率：初始值3e-4，使用余弦退火
- batch size：根据GPU显存尽可能设大，至少32
- 损失权重：检测和分割任务比例建议4:1
- 数据增强：Mosaic+MixUp效果显著，但需注意分割任务的标签同步变换
5.2 常见问题排查
- 显存不足：
- 减小batch size
- 开启梯度累积
- 使用混合精度训练
- 训练不稳定：
- 检查学习率是否过大
- 确认数据标注质量
- 尝试调整损失权重
- 验证指标不提升：
- 检查数据分布是否均衡
- 尝试不同的损失函数组合
- 调整模型容量
5.3 性能优化技巧
- 数据加载优化：
python复制seg_valloader = SegmentationDataset.get_custom_loader(..., workers=8, pin=True)
- 合理设置workers数量（通常为CPU核心数的2-4倍）
- 启用pin_memory加速GPU数据传输
- 训练加速：
- 开启AMP混合精度训练
- 使用DDP分布式训练
- 预加载数据到内存（对小数据集有效）
- 模型部署优化：
- 导出时进行模型剪枝和量化
- 使用TensorRT加速推理
- 针对特定硬件优化算子
6. 扩展与定制
6.1 自定义模型结构
通过修改yaml配置文件，可以灵活调整模型结构：
yaml复制# models/yolov11_custom_seg_big.yaml示例
backbone:
# 主干网络配置
[...]
head:
# 检测头配置
[...]
seg_head:
# 分割头配置
[...]
6.2 支持新任务
框架设计考虑了扩展性，添加新任务的步骤：
- 实现对应的数据加载器
- 设计任务特定的网络头
- 实现相应的损失函数
- 在训练循环中添加任务处理逻辑
6.3 实验管理
框架集成了多种实验管理工具：
- TensorBoard日志记录
- Weights & Biases集成
- 训练过程可视化
python复制if wandb_logger.wandb:
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name)
for x in save_dir.glob('train*.jpg') if x.exists()]})
这套训练框架在实际项目中表现出色，特别是在需要同时处理检测和分割任务的场景下。通过合理的任务权重配置和共享特征设计，可以在不显著增加计算成本的情况下获得多任务输出的优势。对于希望基于YOLOv5进行扩展开发的团队，这个项目提供了很好的起点和参考实现。