YOLOv10 项目 head 模块源码深度解析Detect / Segment / Pose / OBB / Classify / WorldDetect / RTDETRDecoder 全解【免费下载链接】yolov10YOLOv10: Real-Time End-to-End Object Detection [NeurIPS 2024]项目地址: https://gitcode.com/GitHub_Trending/yo/yolov10本篇文章以仓库中 ultralytics/nn/modules/head.py文档见 docs/en/reference/nn/modules/head.md为研究对象系统拆解 YOLOv10 仓库中所有任务头head的实现原理与源码细节。读完本文你将掌握检测、分割、姿态、旋转框、分类、开放词汇检测与 RT-DETR 解码器的完整网络头结构以及它们如何与 YAML 模型配置和训练/推理流程衔接可直接据此进行二次开发、自定义任务头或定位推理问题。一、head 模块在模型中的定位在 YOLO 系列网络结构中模型通常由三部分组成backbone骨干特征提取→ neck多尺度特征融合→ head任务输出。head 位于模型的最末端负责把 neck 输出的多尺度特征图如 P3/P4/P5解码为具体的预测结果——边界框、类别概率、分割掩码、关键点或旋转角等。在 ultralytics/nn/tasks.py 的parse_model()中head 被统一作为cfg[head]列表的最后一个模块实例化elif m in {Detect, WorldDetect, Segment, Pose, OBB, ImagePoolingAttn, v10Detect}: args.append([ch[x] for x in f]) # 把多尺度特征通道列表作为最后一个参数传入 if m is Segment: args[2] make_divisible(min(args[2], max_channels) * width, 8) elif m is RTDETRDecoder: # special case, channels arg must be passed in index 1 args.insert(1, [ch[x] for x in f])从源码结构看所有卷积类任务头Detect 及其子类都继承自同一个nn.Module基类其构造签名统一为(nc, ch)形式——nc是类别数ch是各尺度特征图的通道列表。模型实例化完成后ultralytics/nn/tasks.py 会调用m.stride计算并保存各检测层的下采样步长供 head 后续解码边界框使用。head.py对外导出的模块定义在 ultralytics/nn/modules/head.py__all__ Detect, Segment, Pose, Classify, OBB, RTDETRDecoder下文依次剖析每个类的实现。二、Detect通用锚点检测头所有任务的基类Detect是 YOLOv8/YOLOv10 仓库中最核心的检测头也是Segment、OBB、Pose、WorldDetect、v10Detect的共同基类源码位于 ultralytics/nn/modules/head.py。2.1 构造参数与通道设计class Detect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc # 类别数 self.nl len(ch) # 检测层数一般为 3P3/P4/P5 self.reg_max 16 # DFL 通道数 self.no nc self.reg_max * 4 # 每个锚点输出的通道数 self.stride torch.zeros(self.nl) c2, c3 max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) self.cv2 nn.ModuleList( # 边界框分支 nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch ) self.cv3 nn.ModuleList( # 分类分支 nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch ) self.dfl DFL(self.reg_max) if self.reg_max 1 else nn.Identity()关键设计点nc类别数默认 80COCO。YAML 配置中的nc参数会直接传入构造器。nl检测层层数由传入的通道列表ch长度决定一般为 3对应 P38 倍下采样、P416 倍、P532 倍三个尺度。reg_max16DFLDistribution Focal Loss使用的分布通道数。注释说明ch[0] // 16会在 n/s/m/l/x 不同规格下分别缩放为 4/8/12/16/20即reg_max并非对所有模型都是 16。每个边界框的 4 个边left/top/right/bottom各预测一组 16 维的离散分布因此框分支输出4 * reg_max个通道。no每锚点输出数nc 4 * reg_max是推理时特征图在通道维展开后的总长度。双分支结构cv2框回归分支和cv3分类分支各自由Conv → Conv → nn.Conv2d三层构成输入通道x为各尺度特征图通道输出分别为4*reg_max和nc。2.2 forward 与 training/推理双路径Detect.forward()区分训练与推理两种模式ultralytics/nn/modules/head.pydef forward(self, x): y self.forward_feat(x, self.cv2, self.cv3) # 逐层拼接框分类输出 if self.training: return y # 训练直接返回原始特征 return self.inference(y) # 推理解码为真实框训练模式返回未经解码的原始特征供损失函数如v8DetectionLoss使用推理模式进入inference()完成锚点生成、DFL 解码、坐标缩放与 sigmoid 归一化。inference()的核心逻辑ultralytics/nn/modules/head.py将各尺度输出在通道维拼接为(batch, no, num_anchors)调用make_anchors(x, self.stride, 0.5)生成锚点网格——该函数位于 ultralytics/utils/tal.py以grid_cell_offset0.5网格中心偏移在每一层特征图上生成(h*w, 2)的锚点坐标与对应的步长张量对框分支执行self.dfl(box)离散积分再用dist2bbox把 ltrb 距离转换为 xywh训练时或 xyxy导出时坐标最后乘以self.strides还原到输入图像尺度并对分类分支施加 sigmoid。dist2bbox定义于 ultralytics/utils/tal.py其本质是x1y1 anchor - ltx2y2 anchor rb即锚点分别减去/加上左、上、右、下的预测距离得到完整框的角点。2.3 bias_init训练前的偏置初始化def bias_init(self): for a, b, s in zip(m.cv2, m.cv3, m.stride): a[-1].bias.data[:] 1.0 # box 分支 b[-1].bias.data[: m.nc] math.log(5 / m.nc / (640 / s) ** 2) # cls 分支分类分支的偏置按log(5 / nc / (640/s)²)初始化对应在 640 分辨率下每张图约 0.01 个正样本、共 nc 类的统计先验可显著加快训练早期收敛。该方法在 ultralytics/nn/tasks.py 中被DetectionModel在构建后调用。2.4 DFL 与 Proto两个被复用的基础模块Detect依赖的DFLDistribution Focal Loss 积分模块定义在 ultralytics/nn/modules/block.py。它将(b, 4*reg_max, a)的预测重排为(b, 4, reg_max, a)沿reg_max维做 softmax 后与固定的arange(reg_max)权重做 1×1 卷积从而把离散分布积分成连续的 ltrb 距离——这正是 Generalized Focal Loss 提出的分布表示思想在检测头中的落地。三、Segment实例分割头掩码原型 掩码系数Segment继承自Detectultralytics/nn/modules/head.py在检测分支之外额外输出实例掩码class Segment(Detect): def __init__(self, nc80, nm32, npr256, ch()): super().__init__(nc, ch) self.nm nm # 掩码系数数量 self.npr npr # 原型protos数量 self.proto Proto(ch[0], self.npr, self.nm) # 掩码原型生成器 self.detect Detect.forward c4 max(ch[0] // 4, self.nm) self.cv4 nn.ModuleList( # 掩码系数分支 nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nm, 1)) for x in ch )前向过程分为两条线p self.proto(x[0])由最高分辨率特征图经Proto网络生成(bs, nm, h, w)的掩码原型mc由cv4在三个尺度上输出掩码系数并拼接得到(bs, nm, num_anchors)。推理时最终输出为torch.cat([x[0], mc], 1)——检测结果与掩码系数拼接后处理阶段用系数对原型加权求和mc proto即可还原每个实例的掩码。训练模式下则返回(x, mc, p)三要素供v8SegmentationLoss计算。Proto定义在 ultralytics/nn/modules/block.py由Conv → ConvTranspose2d(2x 上采样) → Conv → Conv组成把输入特征图上采样并压缩到nm个原型通道。常见问题分割模型 YAML如ultralytics/cfg/models/v8/yolov8-seg.yaml中的Segment模块参数为[nc, 32, 256]分别对应上面的nc / nm / npr其中nm32即原型/系数数量。四、OBB旋转框检测头OBBOriented Bounding Box同样继承自Detectultralytics/nn/modules/head.py用于检测带角度的旋转框如遥感 DOTA 数据集场景class OBB(Detect): def __init__(self, nc80, ne1, ch()): super().__init__(nc, ch) self.ne ne # 额外参数数角度默认 1 self.cv4 nn.ModuleList( # 角度分支 nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.ne, 1)) for x in ch )前向时通过cv4预测角度 logits并做(sigmoid() - 0.25) * pi变换把角度映射到[-π/4, 3π/4]区间源码中注释保留了另一种[0, π/2]的映射方案。推理时该角度被存入self.angle属性供重写后的decode_bboxes使用。OBB.decode_bboxes调用dist2rboxultralytics/utils/tal.py——与水平框不同旋转框解码需要先按预测角度做旋转矩阵变换再叠加锚点坐标最终得到旋转框的中心点、宽高和角度。五、Pose关键点检测头Pose继承自Detectultralytics/nn/modules/head.py用于人体姿态估计等关键点任务class Pose(Detect): def __init__(self, nc80, kpt_shape(17, 3), ch()): super().__init__(nc, ch) self.kpt_shape kpt_shape # (关键点数, 每点维度)COCO 默认 (17, 3) self.nk kpt_shape[0] * kpt_shape[1] # 总输出通道数 51 self.cv4 nn.ModuleList( # 关键点分支 nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nk, 1)) for x in ch )kpt_shape决定关键点结构(17, 3)表示 17 个关键点、每点 3 维x, y, visible 可见性(17, 2)则只有 x/y 两维。关键点分支输出(bs, nk, num_anchors)其中通道顺序为x1,y1,v1, x2,y2,v2, ...。关键点解码在kpts_decode()中完成ultralytics/nn/modules/head.pyy[:, 0::ndim] (y[:, 0::ndim] * 2.0 (self.anchors[0] - 0.5)) * self.strides # x y[:, 1::ndim] (y[:, 1::ndim] * 2.0 (self.anchors[1] - 0.5)) * self.strides # y y[:, 2::3] y[:, 2::3].sigmoid() # visiblex/y 坐标解码公式与框解码保持一致的锚点偏移 步长缩放范式可见性通道施加 sigmoid。注意源码中特意以sigmoid()而非sigmoid_()原地操作规避了 Apple MPS 后端的已知 bug。导出 TFLite 时则走另一条显式 reshape 路径以避开 TF 的PLACEHOLDER_FOR_GREATER_OP_CODES问题。六、Classify图像分类头Classify是唯一不继承Detect的轻量分类头ultralytics/nn/modules/head.pyclass Classify(nn.Module): def __init__(self, c1, c2, k1, s1, pNone, g1): super().__init__() c_ 1280 # 隐藏维度参考 efficientnet_b0 self.conv Conv(c1, c_, k, s, p, g) self.pool nn.AdaptiveAvgPool2d(1) self.drop nn.Dropout(p0.0, inplaceTrue) self.linear nn.Linear(c_, c2) def forward(self, x): if isinstance(x, list): x torch.cat(x, 1) x self.linear(self.drop(self.pool(self.conv(x)).flatten(1))) return x if self.training else x.softmax(1)它将输入特征图可为多尺度列表拼接后依次经过Conv → 全局平均池化 → Dropout → Linear把空间信息压缩为c2维类别 logits推理时叠加 softmax 输出概率。c2即类别数nc在 ultralytics/nn/tasks.py 的parse_model()中被特殊处理为与数据集的nc一致。七、WorldDetect开放词汇检测头WorldDetect是 YOLO-World 风格的开放词汇检测头ultralytics/nn/modules/head.py它把文本特征注入检测流程class WorldDetect(Detect): def __init__(self, nc80, embed512, with_bnFalse, ch()): super().__init__(nc, ch) c3 max(ch[0], min(self.nc, 100)) self.cv3 nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, embed, 1)) for x in ch) self.cv4 nn.ModuleList(BNContrastiveHead(embed) if with_bn else ContrastiveHead() for _ in ch)与Detect的关键差异分类分支不再直接输出类别 logits而是把图像区域特征映射到embed维文本嵌入空间再经ContrastiveHead计算区域-文本的相似度得分。前向签名变为forward(self, x, text)——推理时额外接收文本嵌入。支撑模块位于 ultralytics/nn/modules/block.pyContrastiveHead对图像特征和文本特征做 L2 归一化后用einsum(bchw,bkc-bkhw, x, w)计算区域-文本余弦相似度乘以可学习的logit_scale初始化为log(1/0.07)并加上可学习偏置BNContrastiveHead用 BatchNorm 替代 L2 归一化logit_scale初始化为 -1.0源码注释说明 -1.0 更稳定。WorldDetect的推理解码流程与Detect.inference()一致只是类别部分来自对比头输出。八、RTDETRDecoder实时可变形 Transformer 解码器RTDETRDecoderultralytics/nn/modules/head.py是 RT-DETR 模型如ultralytics/cfg/models/rt-detr/rtdetr-l.yaml的 Transformer 解码头与前述卷积头架构完全不同。8.1 构造参数全表参数默认值含义nc80类别数ch(512, 1024, 2048)骨干各层特征通道hd256隐藏层维度nq300query 数量ndp4可变形注意力采样点数nh8多头注意力头数ndl6解码器层数d_ffn1024前馈网络维度dropout0.0Dropout 比率actnn.ReLU()激活函数eval_idx-1评估索引nd100denoising 数量训练去噪label_noise_ratio0.5标签噪声比例box_noise_scale1.0框噪声尺度learnt_init_queryFalse是否学习初始 query 嵌入8.2 内部组成self.input_proj nn.ModuleList( # 骨干特征投影到 hd 维 nn.Sequential(nn.Conv2d(x, hd, 1, biasFalse), nn.BatchNorm2d(hd)) for x in ch) self.decoder DeformableTransformerDecoder(hd, decoder_layer, ndl, eval_idx) self.denoising_class_embed nn.Embedding(nc, hd) # denoising 类别嵌入 self.query_pos_head MLP(4, 2 * hd, hd, num_layers2) # query 位置编码头 self.enc_output nn.Sequential(nn.Linear(hd, hd), nn.LayerNorm(hd)) self.enc_score_head nn.Linear(hd, nc) # 编码器分类头 self.enc_bbox_head MLP(hd, hd, 4, num_layers3) # 编码器框回归头 self.dec_score_head nn.ModuleList([nn.Linear(hd, nc) for _ in range(ndl)]) # 每层解码器分类头 self.dec_bbox_head nn.ModuleList([MLP(hd, hd, 4, num_layers3) for _ in range(ndl)]) # 每层框回归头8.3 前向流程forward(x, batchNone)的四步ultralytics/nn/modules/head.py_get_encoder_input把多尺度特征分别经input_proj投影到hd维展平为(bs, h*w, hd)并拼接get_cdn_group来自 ultralytics/models/utils/ops.py训练时构造 Contrastive DeNoisingCDN组对 GT 框/类别施加label_noise_ratio和box_noise_scale噪声生成 denoising query 与注意力掩码_get_decoder_input_generate_anchors按各层特征图尺寸生成比例锚点宽高随层数指数增长grid_size * 2^i并做 logit 变换经enc_output/enc_score_head打分后以torch.topk选出得分最高的nq个位置作为 querytop-k query selection同时由enc_bbox_head生成动态参考框refer_bboxself.decoderDeformableTransformerDecoder迭代ndl层每层输出经独立的dec_bbox_head/dec_score_head回归框与分类分数返回dec_bboxes, dec_scores以及编码器辅助输出和 denoising 元信息。推理模式下最终输出为torch.cat((dec_bboxes.squeeze(0), dec_scores.squeeze(0).sigmoid()), -1)形状为(bs, 300, 4nc)——300 即nq无需 NMS 即可端到端输出。参数初始化见_reset_parametersultralytics/nn/modules/head.py分类头偏置按bias_init_with_prob(0.01)定义于 ultralytics/nn/modules/utils.py初始化回归头最后一层权重/偏置置零其余用 xavier 初始化。源码注释特别提醒linear_init初始化在自定义数据集训练时可能引发 NaN因此被注释掉。九、YOLOv10 项目特色v10Detect 双分支检测头虽然 API 参考文档列出的导出类是前七节的内容但在本仓库的head.py中还有一个专为 YOLOv10 设计的v10Detectultralytics/nn/modules/head.py它是 YOLOv10项目描述中标注为 NeurIPS 2024 论文端到端、无 NMS能力的核心实现也是 v10 系列 YAML如 ultralytics/cfg/models/v10/yolov10n.yaml中实际使用的 head。class v10Detect(Detect): max_det 300 def __init__(self, nc80, ch()): super().__init__(nc, ch) c3 max(ch[0], min(self.nc, 100)) self.cv3 nn.ModuleList(nn.Sequential( nn.Sequential(Conv(x, x, 3, gx), Conv(x, c3, 1)), # 深度可分离卷积 nn.Sequential(Conv(c3, c3, 3, gc3), Conv(c3, c3, 1)), nn.Conv2d(c3, self.nc, 1)) for i, x in enumerate(ch)) self.one2one_cv2 copy.deepcopy(self.cv2) # one-to-one 框分支 self.one2one_cv3 copy.deepcopy(self.cv3) # one-to-one 分类分支关键设计是双分支并行cv2/cv3one-to-many 分支沿用 DFL 检测头训练时由 one-to-many 匹配提供丰富的正样本监督one2one_cv2/one2one_cv3one-to-one 分支训练时与 one-to-many 分支一起优化推理时仅使用该分支输出配合ops.v10postprocess直接选取 Top-K 结果从而省去 NMS。前向逻辑ultralytics/nn/modules/head.pyone2one self.forward_feat([xi.detach() for xi in x], self.one2one_cv2, self.one2one_cv3) if not self.export: one2many super().forward(x) if not self.training: one2one self.inference(one2one) if not self.export: return {one2many: one2many, one2one: one2one} else: boxes, scores, labels ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det, self.nc) return torch.cat([boxes, scores.unsqueeze(-1), labels.unsqueeze(-1).to(boxes.dtype)], dim-1)值得注意的细节one2one 分支的输入[xi.detach() for xi in x]会切断梯度one2one 分支的训练信号仅来自自身的 one-to-one 匹配损失并采用深度可分离卷积gx/gc3分组卷积降低开销训练与推理均返回字典{one2many: ..., one2one: ...}导出时走v10postprocess位于 ultralytics/utils/ops.py得到(boxes, scores, labels)拼接结果对应的损失函数为v10DetectLoss模型类为YOLOv10DetectionModel两者分别在 ultralytics/nn/tasks.py 与 ultralytics/utils/loss.py 中定义。由于v10Detect继承自DetectDetect.inference()中的 DFL 解码、锚点生成、bias_init等机制在 YOLOv10 中同样生效只是输出侧被替换为双分支结构。十、如何把 head 模块接入自己的模型配置head 模块不是独立运行的组件而是通过 YAML 模型配置在parse_model()中实例化的。以 ultralytics/cfg/models/v8/yolov8.yaml 为例其 head 段最后一行通常为- [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)含义将索引 16、19、22 三层的特征图输入Detect模块nc为类别数。parse_model()会自动把这三层特征图的通道列表ch追加为构造参数。其他任务头对应写法分割Segment, [nc, 32, 256]见 ultralytics/cfg/models/v8/yolov8-seg.yaml姿态Pose, [nc, kpt_shape]见 ultralytics/cfg/models/v8/yolov8-pose.yaml旋转框OBB, [nc]见 ultralytics/cfg/models/v8/yolov8-obb.yaml分类Classify, [nc]见 ultralytics/cfg/models/v8/yolov8-cls.yamlYOLOv10v10Detect, [nc]见 ultralytics/cfg/models/v10/yolov10n.yamlRT-DETRRTDETRDecoder, [nc, 256, 300, 4, 8, 6]见 ultralytics/cfg/models/rt-detr/rtdetr-l.yaml。配置好 YAML 后即可通过标准 API 加载与训练例如from ultralytics import YOLO model YOLO(ultralytics/cfg/models/v10/yolov10n.yaml) # 基于配置构建模型head 的推理输出格式与任务相关检测为(bs, no, num_anchors)的原始特征训练或解码后的框类别推理分割/姿态/旋转框则在检测结果上追加掩码系数/关键点/角度。理解这些输出结构是进行自定义后处理或模型导出的前提。总结与速查Head 类任务额外分支关键依赖Detect水平框检测框DFL 类别make_anchors、dist2bbox、DFLSegment实例分割掩码原型 掩码系数Proto、cv4OBB旋转框检测角度dist2rboxPose关键点关键点坐标/可见性kpts_decodeClassify图像分类全局池化 全连接AdaptiveAvgPool2dWorldDetect开放词汇检测区域-文本对比ContrastiveHeadRTDETRDecoder端到端检测Transformer 解码器DeformableTransformerDecoderv10Detect端到端无 NMS 检测one2one one2many 双分支v10postprocess所有实现均集中在 ultralytics/nn/modules/head.py 这一份文件中配合 ultralytics/nn/modules/block.pyDFL/Proto/对比头、ultralytics/utils/tal.py锚点与框解码工具以及 ultralytics/nn/tasks.py模块装配即可完整理解本项目从配置到推理的全链路。若需扩展自定义任务头推荐以Detect为基类仿照Segment/Pose的模式增加分支并同步实现对应的损失函数与 YAML 配置项。【免费下载链接】yolov10YOLOv10: Real-Time End-to-End Object Detection [NeurIPS 2024]项目地址: https://gitcode.com/GitHub_Trending/yo/yolov10创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考