048、多任务学习扩展:YOLO 同时做检测 + 分割 + 姿态估计的 Head 设计一、从一次“爆显存”的调试说起上个月帮团队调一个多任务模型,需求是:在一张图上同时输出检测框、实例分割mask和17个关键点。我第一反应是“简单,三个head并联嘛”。结果一跑训练,batch size设8,直接OOM。看显存占用,分割head的mask分支吃了将近6个G——因为我把每个检测框都生成了原图尺寸的mask,而不是按框裁剪后的。这个坑让我意识到:多任务head设计,核心不是“怎么加分支”,而是“怎么让分支之间不打架、不浪费”。今天这篇笔记,就围绕YOLO系列做检测+分割+姿态估计的head设计,把踩过的坑和优化思路写清楚。二、多任务head的“三兄弟”架构YOLO的检测head本质是密集预测——在特征图的每个grid cell上预测类别、bbox回归。分割和姿态估计其实也是密集预测,只是输出通道含义不同。2.1 检测head(老本行)# 标准YOLOv8检测head,别动它的结构classDe
048、多任务学习扩展:YOLO 同时做检测 + 分割 + 姿态估计的 Head 设计
048、多任务学习扩展:YOLO 同时做检测 + 分割 + 姿态估计的 Head 设计一、从一次“爆显存”的调试说起上个月帮团队调一个多任务模型,需求是:在一张图上同时输出检测框、实例分割mask和17个关键点。我第一反应是“简单,三个head并联嘛”。结果一跑训练,batch size设8,直接OOM。看显存占用,分割head的mask分支吃了将近6个G——因为我把每个检测框都生成了原图尺寸的mask,而不是按框裁剪后的。这个坑让我意识到:多任务head设计,核心不是“怎么加分支”,而是“怎么让分支之间不打架、不浪费”。今天这篇笔记,就围绕YOLO系列做检测+分割+姿态估计的head设计,把踩过的坑和优化思路写清楚。二、多任务head的“三兄弟”架构YOLO的检测head本质是密集预测——在特征图的每个grid cell上预测类别、bbox回归。分割和姿态估计其实也是密集预测,只是输出通道含义不同。2.1 检测head(老本行)# 标准YOLOv8检测head,别动它的结构classDe