Ultralytics YOLO 模型构建核心剖析:nn/tasks.py 从 YAML 到检测模型的完整流水线

Ultralytics YOLO 模型构建核心剖析:nn/tasks.py 从 YAML 到检测模型的完整流水线 Ultralytics YOLO 模型构建核心剖析nn/tasks.py 从 YAML 到检测模型的完整流水线【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralytics本文以 Ultralytics 仓库中ultralytics/nn/tasks.py的 API 参考页为主体系统讲解该模块承载的三大职责任务模型族的定义DetectionModel、PoseModel、RTDETRDetectionModel、WorldModel 等 12 个模型类、YAML 配置文件到 PyTorch 网络的解析过程parse_model以及权重文件的安全加载与任务推断机制torch_safe_load、load_checkpoint、guess_model_task。读完本文你将能够理解YOLO(yolo26n.yaml)背后模型对象是如何一步步被构建、缩放和加载的并在自定义模型 YAML、排查权重加载报错或扩展新任务模型时找到源码级的依据。对应参考文档见 docs/en/reference/nn/tasks.md核心实现全部位于 ultralytics/nn/tasks.py约 2300 行。1. 模块定位模型层的工厂 加载器ultralytics/nn/tasks.py是 Ultralytics 所有 YOLO 系列模型的统一入口层参考页tasks.md列出的符号可分为三组模型类BaseModel、DetectionModel、OBBModel、SegmentationModel、SemanticSegmentationModel、PoseModel、DepthModel、ClassificationModel、RTDETRDetectionModel、WorldModel、YOLOEModel、YOLOESegModel、Ensemble模型构建函数_initialize_yolo_model、parse_model、yaml_model_load、guess_model_scale权重加载函数temporary_modules、_SafeLoad、torch_safe_load、load_checkpoint、guess_model_task。继承关系上绝大多数任务模型OBB/Segment/Pose/Depth/RTDETR/World/YOLOE都直接继承DetectionModel只覆写init_criterion()或predict()SemanticSegmentationModel与ClassificationModel则直接继承BaseModel。2. BaseModel前向分发、融合与权重迁移BaseModeltasks.py L117是所有模型的公共基类关键行为有四类2.1 forward 的双模分发forward() 用输入类型区分训练与推理if isinstance(x, dict): # 训练/边训边验x 为 {img, cls, bboxes, ...} return self.loss(x, *args, **kwargs) return self.predict(x, *args, **kwargs)而_predict_once()L175-L200实现了按nn.Sequential逐层执行的核心循环每层模块带有i自身序号、f输入来源可为负索引或索引列表、type三个附加属性y缓存中间输出self.save记录哪些层输出必须保留供 Concat 等跨层连接复用embed参数支持在指定层抽取adaptive_avg_pool2d后的嵌入向量用于相似度搜索等场景。2.2 卷积-BatchNorm 融合fuse()L236-L268在导出前被调用将Conv/Conv2/DWConv/ConvTranspose的 BN 折叠进卷积核并对RepConv/RepVGGDW做重参数化对开启 end2end 的Detect头还会额外执行m.fuse()去掉 one2many 分支。is_fused()L270-L280通过统计模型中*Norm*层数量阈值默认 10判断是否已完成融合。2.3 权重加载与跨数据集类别重映射load()L311-L338完成预训练权重到当前模型的迁移先将源权重统一转 FP32 的state_dict调用_remap_cls_by_names()L340-L400——当源/目标模型类别名有交集但数量或顺序不同例如 Objects365 预训练权重微调 COCO时按类别名匹配将源分类头Detect.cv3最后一层的对应行拷贝到目标并从源 dict 中删除这些张量防止intersect_dicts按错误顺序整块拷贝用intersect_dicts求形状一致的交集后以strictFalse载入对首层卷积model.0.conv.weight在输入通道数不同如灰度/多光谱数据时做部分通道拷贝以支持多通道迁移训练最后打印Transferred X/Y items from pretrained weights日志。此外_apply()L292-L309重写了设备/dtype 转换钩子把Detect头中stride、anchors、strides等非参数张量一并搬运避免.to(device)后这些张量滞留在原设备。3. _initialize_yolo_modelYAML 到模型属性的公共装配所有继承DetectionModel的任务模型都通过模块级函数 _initialize_yolo_model() 完成公共初始化步骤为若cfg是路径则经yaml_model_load()解析为 dict并挂到model.yaml向后兼容检测到 YOLOv9 遗留的Silence模块时发出弃用警告并替换为nn.Identity写入输入通道数channels若调用方传入的nc与 YAML 中不同打印Overriding model.yaml nc... with nc...并用参数值覆盖调用parse_model()构建nn.Sequential与save列表初始化占位names{0: 0, 1: 1, ...}与inplace标志。4. 任务模型族从 DetectionModel 到 YOLOESegModel各任务模型的默认配置、损失函数一览默认 cfg 取自各__init__签名损失取自init_criterion实现任务模型类默认 cfg损失函数检测DetectionModelyolo26n.yamlv8DetectionLoss/ end2end 时用E2ELoss旋转框 OBBOBBModelyolo26n-obb.yamlv8OBBLoss/E2ELoss(v8OBBLoss)实例分割SegmentationModelyolo26n-seg.yamlv8SegmentationLoss/E2ELoss(v8SegmentationLoss)语义分割SemanticSegmentationModelyolo26n-sem.yamlSemanticSegmentationLoss姿态估计PoseModelyolo26n-pose.yamlv8PoseLoss或PoseLoss26按头类型深度估计DepthModelyolo26n-depth.yamlDepthLoss26图像分类ClassificationModelyolo26n-cls.yamlv8ClassificationLossRT-DETRRTDETRDetectionModelrtdetr-l.yamlRTDETRDetectionLoss(nc, use_vflTrue)开放词检测WorldModelyolov8s-world.yaml检测损失init_criterion继承YOLOE 检测YOLOEModelyoloe-v8s.yamlTVPDetectLoss视觉提示时/检测损失YOLOE 分割YOLOESegModelyoloe-v8s-seg.yamlTVPSegmentLoss视觉提示时/分割损失4.1 DetectionModelstride 自举与 end2end 模式DetectionModel.__init__()L468-L506在公共装配之外做一件关键的事——自举 stride若末层是Detect则以s 256最小 stride 的 2 倍构造零输入前向一次用256 / 输出空间尺寸反推出每个检测层的 stride 并存入self.stride随后调用一次m.bias_init()。若不是Detect如 RTDETR 解码器stride 取默认值32。该类还提供end2end 属性L508-L529end2end读取末层头的end2end标志YOLO26 默认开启见下文 YAML 示例赋值则通过set_head_attr()下发到头SAHI 式增强推理_predict_augment()L531-L553按scales[1, 0.83, 0.67]、flips[None, 3, None]三组缩放/水平翻转分别前向经_descale_pred()还原坐标、_clip_augmented()裁剪大/小目标的冗余尾部后拼接输出end2end 模型或非Detect头会退回单尺度并打印警告损失初始化init_criterion()L595-L597end2end 模型包装为E2ELoss(self)否则使用v8DetectionLoss(self)。4.2 OBB / Segment / Pose / Depth薄继承 专属损失OBBModelL600、SegmentationModelL632与PoseModelL736、DepthModelL778几乎只覆写init_criterion()。PoseModel额外支持data_kpt_shape参数当数据集 YAML 指定的关键点形状与模型 YAML 不一致时打印Overriding model.yaml kpt_shape... with kpt_shape...并覆盖损失选择依据末层头类型——Pose26头用PoseLoss26否则用v8PoseLoss。DepthModel按类文档说明采用 YOLO backbone FPN DPT 风格稠密深度解码头损失为DepthLoss26。4.3 SemanticSegmentationModel用前向钩子反推最深 strideSemanticSegmentationModelL664-L733直接继承BaseModel。其 stride 自举逻辑值得注意不能只看SemanticSegment头的输入尺寸因为 FPN 会在头之前把 P5 上采样必须跟踪整网前向中最小的 4D 输出空间尺寸来得到最深主干 stride。实现方式是给每层register_forward_hook记录最小高度前向256的零输入后按256 / min_h计算如256/832并在finally中移除钩子。4.4 ClassificationModel独立的_from_yaml与输出层改造分类模型不复用_initialize_yolo_model而是自带_from_yaml()L834-L855通道数优先取 YAML 内channels字段nc既未传参也未在 YAML 中声明时直接抛ValueErrorstride固定为[1]分类没有多尺度约束。静态方法reshape_outputs()L857-L883用于把不同家族的分类头改造成目标类别数YOLO 的Classify头换nn.Linear、ResNet/EfficientNet 等末尾nn.Linear直接替换、nn.Sequential中的nn.Linear或nn.Conv2d按位置替换。4.5 RTDETRDetectionModel查询式解码器的独立 predict/lossRTDETRDetectionModelL890虽继承DetectionModel但主干输出不是网格张量因此重写了两处predict()L1043-L1072逐层前向到self.model[:-1]不含头再把头所需的多尺度特征head([y[j] for j in head.f], batch)一次性送入RTDETRDecoder支持训练时传入batch以做去噪训练denoisingloss()L995-L1041把cls/bboxes/batch_idx整理成 targets并按dn_meta[dn_num_split]把解码输出拆出去噪DN分支与常规分支拼接 encoder 输出后交给RTDETRDetectionLoss注释说明 RTDETR 约有 12 项损失参与反传但日志只展示giou_loss、cls_loss、l1_loss三项_apply()额外搬运解码器的anchors与valid_mask_remap_cls_by_names()被覆写L924-L972RTDETR 的分类张量位于解码器的score_head/class_embed含训练用的denoising_class_embed而非Detect.cv3重映射键集合按此调整。4.6 WorldModel 与 YOLOE文本/视觉提示驱动的开放词模型WorldModelL1075在__init__前先创建txt_feats占位张量与clip_modelNoneset_classes()经build_text_model(clip:ViT-B/32)把类名编码为文本嵌入并同步更新末层nc从而支持不依赖 CLIP 的离线推理。其predict()在逐层前向中区分三类特殊模块C2fAttn接收txt_feats、WorldDetect头接收原始ori_txt_feats、ImagePoolingAttn用图像特征更新txt_feats。YOLOEModelL1202在此之上扩展了get_text_pe()L1241-L1277文本编码器默认取 YAML 中text_model缺省mobileclip:blt输出经head.get_tpe()过辅助文本头get_visual_pe()L1279-L1290把视觉提示喂给模型并以return_vpeTrue取回视觉位置嵌入无提示prompt-free改造set_vocab()L1292-L1322先做一次 warmup 缓存 anchors然后为每个检测层构建LRPCHead并把 cls/loc 分支末层卷积删除完成重参数化nc更新为词表长度get_vocab()L1324-L1349则反向取出融合后的词表层。两者都断言not self.training因为重参数化不可逆predict()L1387-L1428遇到YOLOEDetect头时先算vpe可返回、tpe经get_tpe拼接成cls_pe后作为附加输入x.append(cls_pe)传入头。YOLOESegModelL1455-L1499通过多重继承YOLOEModel SegmentationModel复用检测侧提示逻辑仅在loss()中增加TVPSegmentLoss分支。4.7 EnsembleNMS 式模型集成EnsembleL1502-L1540继承nn.ModuleListforward()对每个成员调用(x, augment, profile)[0]取推理输出然后torch.cat(y, 2)沿通道维拼接源码注释标明这是 nms ensemble即各模型检测框合并后交给后续 NMSmax/mean ensemble 实现保留为注释。5. YAML 到网络parse_model 的解析细节parse_model()L1958-L2266是YAML 字典 →nn.Sequential的核心。以一个真实配置 ultralytics/cfg/models/26/yolo26.yaml 为例其参数区展示了缩放体系nc: 80 end2end: True # whether to use end-to-end mode reg_max: 1 # DFL bins scales: # modelyolo26n.yaml will call yolo26.yaml with scale n # [depth, width, max_channels] n: [0.50, 0.25, 1024] # summary: 260 layers, 2,572,280 parameters, 6.1 GFLOPs s: [0.50, 0.50, 1024] # ... m / l / x backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 3, 2]] # ... - [-1, 2, C2PSA, [1024]] head: # ... - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)parse_model的关键机制顶层参数提取nc / activation / scales / end2end / reg_max / depth_multiple / width_multiple / kpt_shape / scale。若 YAML 带scales而未显式给 scale取第一个键并警告随后按scales[scale]定出depth, width, max_channels激活函数activation字段通过eval(act)重定义Conv.default_act受限加载模式下改用_SafeLoad.activation()的 AST 白名单解析见第 7 节模块解析模块名支持nn.*映射torch.nn、torchvision.ops.*、以及tasks.py全局命名空间中的类Conv、C3k2、Detect等均在 文件头部 L16-L79 从ultralytics.nn.modules导入通道缩放规则对base_modules集合内的层c1 ch[f]、c2 make_divisible(min(c2, max_channels) * width, 8)Classify例外输出必须保持nc重复型模块repeat_modules如C2f、C3k2把n round(n * depth)作为 repeat 参数插入第 2 位C2fAttn额外按width缩放嵌入通道并校验可被nh整除头部参数注入Detect / WorldDetect / YOLOEDetect / Segment / Pose / OBB含 26 系列变体统一追加[reg_max, end2end, [ch[x] for x in f]]三个参数RTDETRDecoder把通道元组插入第 1 位Concat输出通道为各输入通道之和元信息附加每个层挂m_.i序号、m_.f来源、m_.type、m_.np参数量并维护save列表跨层引用需保留输出的层最终返回(nn.Sequential(*layers), sorted(save))。_initialize_yolo_model中parse_model的调用点即 L435。分类 YAML 可对照 ultralytics/cfg/models/26/yolo26-cls.yamlhead 仅一行[-1, 1, Classify, [nc]]且scales的max_channels上限同为 1024/512 档位。6. 配置与元信息查询yaml_model_load / guess_model_scale / guess_model_task6.1 yaml_model_loadyaml_model_load() 解析 YAML 前做两个文件名改写P6 模型yolov8x6形式自动重命名为-p6后缀yolov8x6 → yolov8x-p6并警告规模统一yolov8x.yaml → yolov8.yaml去掉 n/s/m/l/x使同一基座 YAML 配合scale参数复用。加载成功后向 dict 注入scale来自guess_model_scale与yaml_file字段再用check_yaml校验路径存在。6.2 guess_model_scaleguess_model_scale() 用正则yolo(e-)?[v]?\d([nslmx])从文件名中提取规模字母n/s/m/l/x匹配失败返回空字符串。6.3 guess_model_task五级任务推断load_checkpoint在权重未保存 task 字段时调用 guess_model_task()推断优先级为model.args[task]训练时保存的参数model.yaml中 head 末层模块名Classify/cls/fc→classify、含detect→detect、SemanticSegment→semantic、segment→segment、pose→pose、obb→obb、depth→depth遍历modules()按头类型判定SemanticSegment、Segment/YOLOESegment、Classify、Pose、OBB、Depth、Detect/WorldDetect/YOLOEDetect/v10Detect文件路径导出文件ONNX 等内嵌 metadata 中的task以及文件名中的-sem/-seg/-cls/-pose/-obb/-depth或目录名segment/classify/pose/obb/depth/detect全部失败则警告并默认detect提示用户显式指定task。7. 权重加载temporary_modules、_SafeLoad 与 torch_safe_load7.1 temporary_modules旧导入路径兼容temporary_modules() 是一个上下文管理器在with块内把sys.modules[旧路径]指向新模块、把旧模块属性指向新属性退出时清理。torch_safe_load正是用它兼容历史命名空间L1795-L1812modules{ ultralytics.yolo.utils: ultralytics.utils, ultralytics.yolo.v8: ultralytics.models.yolo, ultralytics.yolo.data: ultralytics.data, }, attributes{ ultralytics.nn.modules.block.Silence: torch.nn.Identity, # YOLOv9e ultralytics.nn.tasks.YOLOv10DetectionModel: ultralytics.nn.tasks.DetectionModel, # YOLOv10 ultralytics.utils.loss.v10DetectLoss: ultralytics.utils.loss.E2EDetectLoss, # 跨平台 pathlib pickle 兼容 }7.2 _SafeLoadweights_only 受限加载_SafeLoadL1594-L1763实现可选opt-in的受限检查点加载通过环境变量ULTRALYTICS_SAFE_LOAD定义于 ultralytics/utils/init.py L74SAFE_LOAD env_bool(ULTRALYTICS_SAFE_LOAD)或torch_safe_load(..., safe_onlyTrue)启用。其要点能力探测SUPPORTED hasattr(torch.serialization, get_unsafe_globals_in_checkpoint)需要较新版本 PyTorch旧版自动退化为标准加载允许名单自发现_build()L1684-L1763扫描torch.nn、ultralytics.nn.modules、ultralytics.nn.tasks中所有nn.Module子类连同损失/TAL 类、IterableSimpleNamespace、pathlib跨平台别名、YOLOv9/v10 遗留别名一并注册按 pickle 的module.Name路径建索引按需注册loading()L1617-L1651先用get_unsafe_globals_in_checkpoint扫出该检查点实际引用的全局对象只注册需要的条目源码注释说明允许名单全量注册会显著拖慢反序列化因为受限 unpickler 每条 GLOBAL/REDUCE 指令都要全表重建查找无 eval 的激活解析activation()L1653-L1681受限模式下parse_model解析activation字段改用 AST 解析仅接受[torch.]nn.Module(字面量参数)形式其余一律拒绝。7.3 torch_safe_load容错与诊断torch_safe_load() 在别名上下文与可选的受限加载上下文中执行torch.load(file, map_locationcpu)并把各类失败映射为可操作的报错TorchScript 归档RuntimeError含 TorchScript archive 或加载结果是ScriptModule时抛TypeError提示导出 torchscript 或直接加载原始.pt文件损坏截断 zipRuntimeError/PytorchStreamReader、空文件EOFError、非 pickle 字节UnpicklingError统一为 文件为空/截断/损坏 的报错若只是官方资产按裸名请求且缓存损坏会自动删除缓存重新下载用户自定义路径则直接报错YOLOv5 权重缺失模块为models / models.yolo / models.common / models.experimental时明确说明YOLOv5 模型与 ultralyticsv8不向前兼容建议重新训练或使用官方模型缺失依赖numpy._core缺失提示 numpy1.26.1 版本要求其他缺失模块在默认路径下触发check_requirements自动安装受限加载路径下禁止自动安装格式校验加载结果不是 dict例如torch.save(model, ...)保存的实例时警告并按{model: ckpt.model}包装。7.4 load_checkpoint装配推理态模型load_checkpoint() 完成最终装配远程前缀如 http先经check_file下载到weights_dir模型候选优先ckpt[ema]否则ckpt[model]且必须是nn.Module否则报非官方检查点格式错误随后统一 FP32挂model.args默认配置与训练参数合并、以模型参数优先与model.pt_path缺失时补model.task guess_model_task(model)、model.stride [32.0]按需fuse()后置eval()并移到目标设备最后遍历模块同步inplace标志、修补旧版nn.Upsample.recompute_scale_factor。8. 组合示例一条训练/推理链路如何经过本模块从源码结构看典型调用链为上层YOLO引擎 →load_checkpoint()YOLO(yolo26n.pt)时走torch_safe_load→ EMA 权重 →guess_model_task→eval而YOLO(yolo26n.yaml)训练初始化时则直接实例化DetectionModel(cfg)→_initialize_yolo_model→yaml_model_load注入scale→parse_model按scales.n缩放到 260 层/2.57M 参数→DetectionModel.__init__的 256 零输入前向自举 stride 并initialize_weights。训练每步model(batch_dict)经forward → loss → init_criterionYOLO26 默认 end2end即E2ELoss包装导出前再经fuse()折叠 BN 完成重参数化。分类任务链路则走ClassificationModel._from_yaml → v8ClassificationLoss。9. 小结与延伸阅读ultralytics/nn/tasks.py以约 2300 行代码集中了 Ultralytics 模型层的三大问题如何按任务族复用同一套构建逻辑BaseModel_initialize_yolo_model 薄继承、如何把声明式 YAML 确定性翻译成可训练网络parse_model的通道/深度缩放与头部参数注入、如何安全且可诊断地加载历史与第三方权重temporary_modules别名、_SafeLoad白名单、torch_safe_load错误分类。扩展新任务模型时参照OBBModel/DepthModel的最小覆写模式__init__指定默认 YAML init_criterion指定损失即可接入整个引擎而排查加载 .pt 报错时第 7.3 节的错误映射表几乎覆盖了所有官方模型场景。相关文档与实现参考页、YOLO26 检测 YAML、YOLO26 分类 YAML、nn/modules 模块库、自动后端加载。【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralytics创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考