Ultralytics YOLO模块自定义与替换听起来像是算法工程师的进阶课程。实际上我这两年被问得最多的问题反而不是训练怎么跑而是“我想把Backbone换了怎么办”“损失函数能不能自己改”“换了模块之后报错怎么查”。这篇文章就是把我在真实项目里动刀YOLO源码的经验拿出来讲透包括改yaml、写子类、甚至直接改损失函数适合已经跑通过一次yolo训练、想进一步做模型改进的读者。我会带着你从源码结构开始把三种常用的自定义与替换路径讲清楚再用两个完整实战案例演示替换Backbone和损失函数最后把踩过的坑和排查思路一并列出来。1. 先把话说清楚市面上那么多YOLO工具为什么我还在改Ultralytics源码1.1 默认模型满足不了的需求才是自定义的起点很多人一开始用Ultralytics YOLO都是装好包、跑通train觉得一切都很顺。直到某一刻发现自己的数据集特别难搞或者部署环境对体积和帧率有硬性要求。我在一个工业表面缺陷检测项目里就遇到了这样的事默认的YOLOv8l模型准确率不错但嵌入式设备跑起来只有8帧/s必须把模型压到很小。这时候绕不开的问题就出现了能不能把C2f替换成更轻量的结构能不能拿自己训练的Backbone直接接上YOLO的检测头这就是“Ultralytics YOLO模块自定义与替换”最典型的场景。默认模型不是万能的数据集尺度分布、目标形态、算力平台都会逼着你改结构。另一个常见需求是改进实验很多论文里的改进点比如注意力模块、可变形卷积、新的特征融合方式都需要以“新增模块”的形式注入到YOLO里。如果只会用现成模型论文就复现不了自己的思路也验证不了。1.2 三种自定义方式的范围与成本对比我在项目里总结下来模块自定义与替换大致有三条路成本从低到高方式改动范围适合场景维护成本风险改yaml配置只调整已存在模块的组合与参数快速对比结构、改深度宽度、换现有模块名低低写子类并注册新增Python网络模块注入给YOLO使用实现新注意力、新Block、新Backbone中中直接改源码修改任务逻辑、损失函数、前向流程深度定制训练流程、自定义损失、改动检测头逻辑高高升级极易冲突需要说清楚的是这三种路径不是互斥的。我通常先改yaml做快速验证确认方向后把稳定模块整理成子类放进单独包管理。直接改源码则是万不得已的选择比如要替换官方损失函数里某个算子短期内又不想维护自己的分支只能在fork里改但升级时一定要小心。2. 动刀之前先摸清Ultralytics YOLO的目录结构和模块加载机制2.1 ultralytics/nn模块到底放了哪些东西要“替换”模块就得先知道模块从哪来。Ultralytics YOLO和早期版本的Darknet框架不一样它整个项目是纯PyTorch代码核心目录很清晰ultralytics/ nn/ backbone.py # 经典Backbone实现比如某些预训练模型兼容层 modules.py # 所有基本模块Conv, C2f, SPPF, Detect, DFL等 tasks.py # 模型构建、加载、前向逻辑parse_model在这里 autobackend.py # 推理时自动选择后端PyTorch、ONNX、TensorRT等 cfg/ models/ v8/ yolov8.yaml yolov8-seg.yaml yolov8-pose.yaml v11/ yolov11.yaml yolov11-seg.yaml其中modules.py是“零件库”所有可复用的卷积、残差块、空间金字塔池化、检测头都在这里定义。tasks.py是“组装车间”它读取yaml结构描述在parse_model函数中把零件一个个拼起来。如果你想加一个全新模块通常就是先在modules.py里写类然后让tasks.py能按名字找到它。2.2 模型yaml文件与Python模块的映射关系YOLOv8的yolov8.yaml长这样里面没有可训练参数只描述网络结构backbone: - [-1, 1, Conv, [64, 3, 2]] - [-1, 1, Conv, [128, 3, 2]] - [-1, 1, C2f, [128, True]] - [-1, 1, Conv, [256, 3, 2]] - [-1, 1, C2f, [256, True]] ... head: - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 6], 1, Concat, [1]] - [-1, 1, C2f, [128]] ...每一行的意思输入来自哪一层-1表示上一层有几个模块模块类名以及构造参数列表。tasks.py在解析时会把类名字符串映射到实际类。注意它不仅能映射ultralytics.nn.modules里的类还能映射torch下的基础模块比如nn.Upsample因为代码里会遍历一个包含nn.modules内容的字典。2.3 模块注册表是怎么工作的很多初学者直接在自己写的my_module.py里定义好类然后在yaml里写MyModule结果一加载就报KeyError: MyModule。原因是Ultralytics没有像某些框架那样用装饰器做全局注册表它是靠tasks.py里的导入和__all__来控制可见性的。如果你想新增模块最直接的做法是在ultralytics/nn/modules.py里新增类。在ultralytics/nn/modules/__init__.py的__all__中导出该类。确认tasks.py里已经from ultralytics.nn.modules import *这样解析时就能通过类名字符串找到它。我不建议把自定义模块直接塞进modules.py主文件里因为升级时会很难看diff。更优雅的做法是把自定义模块放到你自己的项目包中然后在tasks.py里手动导入并加入一个映射字典。虽然要多写几行但升级时只需要处理一个patch点。3. 三种自定义与替换路径改配置、写子类、直接动源码3.1 只改yaml配置就能完成的替换如果你的目标是调整整体结构而不是实现一个全新的算子那么只改yaml可能是最省事的。举个例子YOLOv11新增了C2PSA模块如果你想把v8结构里的C2f替换成C2PSA只要把yaml里所有C2f改成C2PSA并且保证参数形式兼容。如# 原始 - [-1, 1, C2f, [512, True]] # 替换后 - [-1, 1, C2PSA, [512, True]]这种“字符串替换”的方式虽然看着简单但前提是模块的构造函数参数个数和含义要能接得上。很多人在这一步摔跟头是因为只改了名称没有检查后面参数列表。比如C2PSA和C2f的通道数、是否使用残差的语义并不完全一致直接替换后可能不报错但训练效果会诡异。所以在做配置级替换时务必先用model.info()看一下每一层的输出shape再做训练验证。3.2 通过继承与注册写自定义模块要实现论文里的新结构比如一个带SE注意力的Bottleneck就得走“写子类”的路线。这里给一个真实用过的例子import torch import torch.nn as nn from ultralytics.nn.modules import Conv class SEBottleneck(nn.Module): 带Squeeze-and-Excitation的残差块可替换YOLO中的Bottleneck def __init__(self, c1, c2, shortcutTrue, g1, k(1, 3), e0.5): super().__init__() c_ int(c2 * e) # hidden channels self.cv1 Conv(c1, c_, k[0], 1) self.cv2 Conv(c_, c2, k[1], 1, gg) self.add shortcut and c1 c2 # SE模块 self.se nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(c2, c2 // 4, kernel_size1), nn.ReLU(inplaceTrue), nn.Conv2d(c2 // 4, c2, kernel_size1), nn.Sigmoid() ) def forward(self, x): y self.cv2(self.cv1(x)) y y * self.se(y) return x y if self.add else y用这个模块替换yaml中的某个Bottleneck核心步骤是将上述代码放到你的自定义模块文件比如my_modules.py里并保证from ultralytics.nn.modules import Conv能正常工作。在tasks.py里导入from my_modules import SEBottleneck。在parse_model的模块映射字典中加上SEBottleneck或者把它导出到ultralytics.nn.modules的命名空间。修改你的模型yaml把某个Bottleneck改成SEBottleneck。这样做的好处是基本不影响原包结构坏处是升级Ultralytics版本时需要重新添加映射。我一般会在自己项目里写一个register.py专门集中处理这些导入和映射方便统一改。3.3 直接修改源码的利与弊第三种方式最直接但也最需要储备。比如你想自定义损失函数只改yaml和新增子类都没用因为损失函数是在ultralytics/utils/loss.py里被硬编码调用的。这时候你只有两条路一是fork官方仓库在loss.py里改二是通过monkey-patch在训练脚本里覆盖原有类。我在早期项目里试过直接改源码确实方便当时想加一个辅助分类头直接在DetectionModel.forward里加了个分支很快验证了想法。但发版后遇到bug想同步上游更新时git冲突多到怀疑人生。后来我学乖了如果能用子类继承、注册覆盖的方式就不要改原文件。如果实在要改就把改动的部分抽成独立函数在源码里只留一行调用这样diff最小。4. 实战一自研轻量Backbone替换YOLOv8默认Backbone4.1 目标场景与设计思路嵌入式设备上跑缺陷检测输入图像不大但对延迟敏感。默认YOLOv8s的Backbone包含多次下采样和大量C2f对算力要求还是偏高。我的设计目标是做一个三层下采样的轻量Backbone保持特征提取能力的同时减少通道数。结构上我选择了“卷积倒残差轻量注意力”的路线。虽然Ultralytics自带了一些经典Backbone但为了让读者明白完整的替换过程我以一个自定义LightBackbone为例避免你只会改官方已有模块、遇到新结构却不会注入。4.2 编写自定义模块并注册新建light_backbone.py定义整个Backbone结构。这里不用太复杂突出“能被yaml直接实例化”的关键点import torch import torch.nn as nn from ultralytics.nn.modules import Conv class LightStem(nn.Module): 下采样stem使用两个conv快速降低分辨率 def __init__(self, c1, c2): super().__init__() self.conv1 Conv(c1, c2, k3, s2) self.conv2 Conv(c2, c2, k3, s1) def forward(self, x): return self.conv2(self.conv1(x)) class LightBlock(nn.Module): 轻量特征提取块通道减半使用深度可分离卷积 def __init__(self, c1, c2, n1, shortcutTrue): super().__init__() cv1 Conv(c1, c2, 1, 1) cv2 Conv(c2, c2, 3, 1, gc2) # depthwise cv3 Conv(c2, c2, 1, 1) self.block nn.Sequential(cv1, cv2, cv3) self.add shortcut and c1 c2 def forward(self, x): return x self.block(x) if self.add else self.block(x) class LightBackbone(nn.Module): def __init__(self, c13, c264): super().__init__() self.stem LightStem(c1, c2) self.stage1 LightBlock(c2, c2, n1) self.down2 Conv(c2, c2 * 2, 3, 2) self.stage2 LightBlock(c2 * 2, c2 * 2, n2) self.down3 Conv(c2 * 2, c2 * 4, 3, 2) self.stage3 LightBlock(c2 * 4, c2 * 4, n2) self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(c2 * 4, c2 * 4) # 仅用于输出特征实际可去掉 def forward(self, x): x self.stem(x) x self.stage1(x) x self.down2(x) x self.stage2(x) x self.down3(x) x self.stage3(x) return x注意这个自定义Backbone需要输出YOLO检测头需要的特征图列表。为了简化这里只演示注册方法实际替换Backbone时你要让不同尺度的特征图都能从Backbone中返回。我的做法是在LightBackbone.forward里返回[x_stage1, x_stage2, x_stage3]三个特征图分别对应80×80、40×40、20×20的检测尺度。然后把它注册到Ultralytics中。在tasks.py头部加入from light_backbone import LightBackbone再在parse_model可用的模块字典中补上m eval_name if m in {...}: ... else: m globals()[m]最省事的注册方式是直接把类导入到ultralytics.nn.modules.__init__.py的__all__里让tasks.py的from ultralytics.nn.modules import *自动拿到类名。我在实际项目中更愿意维护一个patching.py在训练脚本启动时执行tasks.parse_model里的字典更新避免改动包本体。4.3 修改模型yaml并验证结构现在写一个yolov8-light.yaml只保留YOLOv8的检测头替换掉整段Backbone# Ultralytics YOLOv8-light.yaml backbone: - [ -1, 1, LightBackbone, [3, 64] ] head: - [ -1, 1, Conv, [ 128, 3, 2 ] ] - [ -1, 1, C2f, [ 128, True ] ] - [ -1, 1, SPPF, [ 128, 5 ] ] - [ -1, 1, Conv, [ 64, 1, 1 ] ] - [ -1, 1, nn.Upsample, [ None, 2, nearest ] ] - [ [ -1, 3 ], 1, Concat, [ 1 ] ] ...这里做了简化真正的head需要对应Backbone输出的特征图索引。我把Backbone整体当作v8原结构中第1到第9层的替代检测头里的Concat层引用了Backbone内部输出因此在yaml中要通过多个层号串联。这个最容易错我的建议是先用一个自己写的build_model.py打印结构不要急着训练from ultralytics import YOLO model YOLO(yolov8-light.yaml) model.info()运行后如果能看到每一层的输出shape并且没有KeyError说明模块注册和yaml结构基本正确。如果model.info()报RuntimeError通常是你自定义模块的forward返回的特征图数量、通道与检测头不匹配。4.4 训练与性能对比结构验证通过后用小型数据集先跑10个epoch我一般会在train命令里加上--workers 4 --batch 16 --epochs 10观察loss是否能下降。对比指标时除了mAP50、mAP50-95还要看参数量、GFLOPs和单帧耗时。我在那个缺陷项目里换了LightBackbone后参数量降到原来的38%帧率从8帧提升到25帧代价是mAP50下降了1.2%。如果你的任务对准确率更敏感就不要盲目追求轻量化可以尝试在Block里加SE或注意力再对比。5. 实战二替换损失函数解决小目标和样本不均衡问题5.1 YOLO默认损失函数的结构很多人在自定义模块时只关注网络结构忽略了损失函数才是影响训练结果的关键。YOLOv8的默认损失在ultralytics/utils/loss.py中核心是v8DetectionLoss它由三部分组成BoxLoss边界框损失通常是CIoU加上DFLDistribution Focal Loss。ClsLoss分类损失默认二值交叉熵BCE。DFL分布焦点损失用于回归框的分布表示。如果数据集存在严重的类别不均衡或者大量小目标默认BCE分类损失往往不够用。这时候就需要替换损失函数比如换成Focal Loss。5.2 自定义FocalLoss并替换在loss.py中v8DetectionLoss的__init__里会执行self.bce nn.BCEWithLogitsLoss(reductionnone)如果你想替换成Focal Loss最简单的做法是写一个FocalLoss类然后通过继承原类覆盖bce属性import torch import torch.nn as nn class FocalLoss(nn.Module): def __init__(self, gamma1.5, alpha0.25): super().__init__() self.gamma gamma self.alpha alpha def forward(self, pred, target): pred_sigmoid pred.sigmoid() pt (1 - pred_sigmoid) * target pred_sigmoid * (1 - target) focal_weight (1 - pt).pow(self.gamma) if self.alpha is not None: alpha_t self.alpha * target (1 - self.alpha) * (1 - target) focal_weight focal_weight * alpha_t loss nn.functional.binary_cross_entropy_with_logits( pred, target, reductionnone ) return focal_weight * loss然后在训练脚本里继承官方损失类并替换from ultralytics.utils.loss import v8DetectionLoss class CustomV8Loss(v8DetectionLoss): def __init__(self, model): super().__init__(model) self.bce FocalLoss(gamma1.5)最后通过custom训练入口使用这个损失。需要注意v8DetectionLoss的__init__可能已经在后续版本中调整了bce的作用范围比如用在分类损失处也可能用于DFL因此要看清代码。更稳妥的替换点是把v8DetectionLoss里的self.bce换成self.bce FocalLoss(...)但不能影响DFL内部的BCE使用因为DFL的数学形式和分类不完全一样。我在v8.2版本上实际验证过直接替换self.bce后DFL也变成Focal形式训练loss曲线变得很奇怪。后来我单独保留了一个self.bce_cls给分类分支才解决了问题。所以替换损失函数时一定要理清代码结构不能一刀切。5.3 训练中观察loss曲线与指标变化替换损失函数后不要只看mAP。先用plotsTrue训练观察train/cls_loss、train/box_loss、train/dfl_loss三条曲线。如果分类loss明显下降变缓说明Focal Loss的gamma偏大。如果目标框回归出现问题则可能是误替换了DFL部分。我会在验证集上额外统计各类别的AP而不是只看平均mAP这样才能确认是否真的解决了小目标类别的不均衡问题。6. 替换Backbone/Head必踩的坑从报错信息到修复全链路6.1 常见报错与根因对照我整理了替换模块时最常见的报错按出现频率排序报错信息根因解决方案KeyError: MyModule自定义类没有注册到tasks.py的解析环境导入模块并在模块字典中添加名字RuntimeError: size mismatch特征图通道数或层索引不匹配打印每一层输出shape检查yaml中的层索引引用TypeError: __init__() got an unexpected keyword argumentyaml中参数列表顺序或名称与构造函数不一致对照构造函数签名调整yaml参数CUDA out of memory替换后的模块显存占用更大降低batch size用--device cpu先验证结构再上GPUAttributeError: NoneType object has no attribute shape自定义模块在某个分支上返回了None检查forward路径避免条件分支缺失6.2 一次维度不匹配的完整排查记录有一次我把新设计的Neck层接进YOLOv8-seg训练到第一个step就报size mismatch。我没急着看堆栈而是先把模型切到CPU用随机张量跑了一次前向import torch model YOLO(my-seg.yaml).model x torch.randn(1, 3, 640, 640) with torch.no_grad(): y model(x)结果提示torch.cat时形状不一致。我就在tasks.py的parse_model后面临时加了打印把每一层输出的shape打出来。最后发现是yaml里Concat层引用的是第3层输出而第3层经过我替换的Block后通道数是128预期是192。问题不是模块本身而是我在yaml的层索引写错了。修正索引后一切正常。这个经历说明自定义模块后第一件事永远是“打印结构、单步前向”而不是直接扔给训练。6.3 设备与精度相关的坑AMD显卡、半精度等Ultralytics官方现在对AMD显卡的支持主要靠ROCm和DirectML。我在AMD RX6700上跑自定义模块时遇到过torch.cat算子在ROCm后端触发异常换成CPU却没有问题。这类算子兼容性坑很隐蔽建议先用小batch在CPU上验证一遍网络前向再换AMD显卡。另外开启AMP混合精度训练后自定义模块里的某些激活函数或自定义op可能不在autocast支持列表中导致loss变成NaN。最有效的排查方法是关闭AMP也就是训练时加--amp False如果loss恢复正常那就是精度不兼容。处理方式是在自定义模块的关键计算前后手动加torch.cuda.amp.autocast(enabledFalse)或者重写其中的算子。7. 进阶玩法数据集、实例分割和多目标跟踪模块的替换与适配7.1 替换数据加载与增强模块适配自己的数据集模块替换不只是网络结构的事。很多人在训练自己数据集时发现acc上不去问题常常出在数据加载环节。Ultralytics的BaseDataset和augment.py支持自定义增强器。我习惯用Albumentations库替换部分内置增强做法是继承BaseDataset在get_image_and_label里调用自己的增强管线。替换数据增强模块时要特别注意增强后需要同步更新bbox坐标比如随机旋转、随机裁剪都会改变框坐标。如果不小心只增强图像不改框模型就会学到错误映射。7.2 实例分割任务中的检测头替换YOLOv8-seg和YOLOv11-seg的检测头是Segment类它在普通Detect头基础上增加了一个掩码分支。如果你想把检测头替换成自定义分割头需要同时修改yaml的head部分和tasks.py里的Segment组装逻辑。分割头通常还要输出原型掩码系数和检测分支的边界框特征结合。很多人在替换分割头时只改了nn.Module却忘记修改v8SegmentationLoss里的掩码损失计算方式结果训练时mask_loss恒为0。我在项目里通常保留原损失函数先替换主干网络等效果稳定再动分割头。7.3 多目标跟踪模块的接入与替换Ultralytics的预测接口里集成了ByteTrack、BoT-SORT等跟踪器。如果你有自研的跟踪器可以在ultralytics/trackers/track.py中替换on_predict_start中的tracker实例。常见做法是继承BaseTracker在update方法中接受检测结果和嵌入特征返回跟踪ID。替换跟踪模块时最难处理的是“检测框和ReID特征如何对齐”因为YOLO的推理结果可能包含多个batch。我的经验是先关闭batch即predict(source..., batch1)能省掉大量调试时间。8. 版本升级时模块变动的影响与兼容性处理8.1 v8到v11哪些模块API变了Ultralytics版本迭代很快从v8到v11yaml里的模块名和代码结构都有变化。比如v11引入了C3k2、C2PSA部分旧模块虽然保留但不再是默认推荐。tasks.py解析逻辑也在变化之前使用的model.model[-1]等索引方式可能会失效。如果你用的是几个月前的自定义模块升级后很可能出现两类问题一是模块名被移除二是构造函数参数比如增加了一个c3参数导致yaml参数列表不兼容。我的建议是升级前先看官方仓库的Release Notes重点看nn/modules.py和loss.py的diff。不要直接升级大版本先在虚拟环境里跑一次model.info()和单step训练确认自定义模块的状态。8.2 自定义模块跨版本迁移方法把自定义模块从源码中抽离出来是降低升级痛苦最有效的手段。我在自己的项目里维护了一个extension/目录里面每个模块都是独立的nn.Module并且提供一个register()函数负责在Ultralytics的模块映射表中插入自定义类名。迁移时只需要做三件事确认新版本的目标类构造函数签名是否变化。修改register()里的映射关系。运行全部训练和推理测试。只要你没有改动官方源码升级时直接从Git拉取新版本即可。如果确实需要改官方源码我强烈建议用git format-patch保存一个独立补丁升级时git apply后再清理冲突而不是直接把所有改动提交到官方仓库分支里。我在多个项目里用过Ultralytics YOLO之后最大的体会是模块自定义与替换并不难难的是搞清楚它背后的加载机制和参数约定。只要你愿意花半天时间读完tasks.py里的parse_model函数再动手写第一个自定义模块后面所有替换工作都会变得非常顺手。另外自定义模块一定要写在独立文件里而不是直接改modules.py否则版本升级时会特别痛苦。如果你正准备给YOLO换Backbone或调损失函数不妨从上面的一个小案例开始先跑通再深入。
Ultralytics YOLO模块自定义与替换：从原理到实战
Ultralytics YOLO模块自定义与替换听起来像是算法工程师的进阶课程。实际上我这两年被问得最多的问题反而不是训练怎么跑而是“我想把Backbone换了怎么办”“损失函数能不能自己改”“换了模块之后报错怎么查”。这篇文章就是把我在真实项目里动刀YOLO源码的经验拿出来讲透包括改yaml、写子类、甚至直接改损失函数适合已经跑通过一次yolo训练、想进一步做模型改进的读者。我会带着你从源码结构开始把三种常用的自定义与替换路径讲清楚再用两个完整实战案例演示替换Backbone和损失函数最后把踩过的坑和排查思路一并列出来。1. 先把话说清楚市面上那么多YOLO工具为什么我还在改Ultralytics源码1.1 默认模型满足不了的需求才是自定义的起点很多人一开始用Ultralytics YOLO都是装好包、跑通train觉得一切都很顺。直到某一刻发现自己的数据集特别难搞或者部署环境对体积和帧率有硬性要求。我在一个工业表面缺陷检测项目里就遇到了这样的事默认的YOLOv8l模型准确率不错但嵌入式设备跑起来只有8帧/s必须把模型压到很小。这时候绕不开的问题就出现了能不能把C2f替换成更轻量的结构能不能拿自己训练的Backbone直接接上YOLO的检测头这就是“Ultralytics YOLO模块自定义与替换”最典型的场景。默认模型不是万能的数据集尺度分布、目标形态、算力平台都会逼着你改结构。另一个常见需求是改进实验很多论文里的改进点比如注意力模块、可变形卷积、新的特征融合方式都需要以“新增模块”的形式注入到YOLO里。如果只会用现成模型论文就复现不了自己的思路也验证不了。1.2 三种自定义方式的范围与成本对比我在项目里总结下来模块自定义与替换大致有三条路成本从低到高方式改动范围适合场景维护成本风险改yaml配置只调整已存在模块的组合与参数快速对比结构、改深度宽度、换现有模块名低低写子类并注册新增Python网络模块注入给YOLO使用实现新注意力、新Block、新Backbone中中直接改源码修改任务逻辑、损失函数、前向流程深度定制训练流程、自定义损失、改动检测头逻辑高高升级极易冲突需要说清楚的是这三种路径不是互斥的。我通常先改yaml做快速验证确认方向后把稳定模块整理成子类放进单独包管理。直接改源码则是万不得已的选择比如要替换官方损失函数里某个算子短期内又不想维护自己的分支只能在fork里改但升级时一定要小心。2. 动刀之前先摸清Ultralytics YOLO的目录结构和模块加载机制2.1 ultralytics/nn模块到底放了哪些东西要“替换”模块就得先知道模块从哪来。Ultralytics YOLO和早期版本的Darknet框架不一样它整个项目是纯PyTorch代码核心目录很清晰ultralytics/ nn/ backbone.py # 经典Backbone实现比如某些预训练模型兼容层 modules.py # 所有基本模块Conv, C2f, SPPF, Detect, DFL等 tasks.py # 模型构建、加载、前向逻辑parse_model在这里 autobackend.py # 推理时自动选择后端PyTorch、ONNX、TensorRT等 cfg/ models/ v8/ yolov8.yaml yolov8-seg.yaml yolov8-pose.yaml v11/ yolov11.yaml yolov11-seg.yaml其中modules.py是“零件库”所有可复用的卷积、残差块、空间金字塔池化、检测头都在这里定义。tasks.py是“组装车间”它读取yaml结构描述在parse_model函数中把零件一个个拼起来。如果你想加一个全新模块通常就是先在modules.py里写类然后让tasks.py能按名字找到它。2.2 模型yaml文件与Python模块的映射关系YOLOv8的yolov8.yaml长这样里面没有可训练参数只描述网络结构backbone: - [-1, 1, Conv, [64, 3, 2]] - [-1, 1, Conv, [128, 3, 2]] - [-1, 1, C2f, [128, True]] - [-1, 1, Conv, [256, 3, 2]] - [-1, 1, C2f, [256, True]] ... head: - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 6], 1, Concat, [1]] - [-1, 1, C2f, [128]] ...每一行的意思输入来自哪一层-1表示上一层有几个模块模块类名以及构造参数列表。tasks.py在解析时会把类名字符串映射到实际类。注意它不仅能映射ultralytics.nn.modules里的类还能映射torch下的基础模块比如nn.Upsample因为代码里会遍历一个包含nn.modules内容的字典。2.3 模块注册表是怎么工作的很多初学者直接在自己写的my_module.py里定义好类然后在yaml里写MyModule结果一加载就报KeyError: MyModule。原因是Ultralytics没有像某些框架那样用装饰器做全局注册表它是靠tasks.py里的导入和__all__来控制可见性的。如果你想新增模块最直接的做法是在ultralytics/nn/modules.py里新增类。在ultralytics/nn/modules/__init__.py的__all__中导出该类。确认tasks.py里已经from ultralytics.nn.modules import *这样解析时就能通过类名字符串找到它。我不建议把自定义模块直接塞进modules.py主文件里因为升级时会很难看diff。更优雅的做法是把自定义模块放到你自己的项目包中然后在tasks.py里手动导入并加入一个映射字典。虽然要多写几行但升级时只需要处理一个patch点。3. 三种自定义与替换路径改配置、写子类、直接动源码3.1 只改yaml配置就能完成的替换如果你的目标是调整整体结构而不是实现一个全新的算子那么只改yaml可能是最省事的。举个例子YOLOv11新增了C2PSA模块如果你想把v8结构里的C2f替换成C2PSA只要把yaml里所有C2f改成C2PSA并且保证参数形式兼容。如# 原始 - [-1, 1, C2f, [512, True]] # 替换后 - [-1, 1, C2PSA, [512, True]]这种“字符串替换”的方式虽然看着简单但前提是模块的构造函数参数个数和含义要能接得上。很多人在这一步摔跟头是因为只改了名称没有检查后面参数列表。比如C2PSA和C2f的通道数、是否使用残差的语义并不完全一致直接替换后可能不报错但训练效果会诡异。所以在做配置级替换时务必先用model.info()看一下每一层的输出shape再做训练验证。3.2 通过继承与注册写自定义模块要实现论文里的新结构比如一个带SE注意力的Bottleneck就得走“写子类”的路线。这里给一个真实用过的例子import torch import torch.nn as nn from ultralytics.nn.modules import Conv class SEBottleneck(nn.Module): 带Squeeze-and-Excitation的残差块可替换YOLO中的Bottleneck def __init__(self, c1, c2, shortcutTrue, g1, k(1, 3), e0.5): super().__init__() c_ int(c2 * e) # hidden channels self.cv1 Conv(c1, c_, k[0], 1) self.cv2 Conv(c_, c2, k[1], 1, gg) self.add shortcut and c1 c2 # SE模块 self.se nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(c2, c2 // 4, kernel_size1), nn.ReLU(inplaceTrue), nn.Conv2d(c2 // 4, c2, kernel_size1), nn.Sigmoid() ) def forward(self, x): y self.cv2(self.cv1(x)) y y * self.se(y) return x y if self.add else y用这个模块替换yaml中的某个Bottleneck核心步骤是将上述代码放到你的自定义模块文件比如my_modules.py里并保证from ultralytics.nn.modules import Conv能正常工作。在tasks.py里导入from my_modules import SEBottleneck。在parse_model的模块映射字典中加上SEBottleneck或者把它导出到ultralytics.nn.modules的命名空间。修改你的模型yaml把某个Bottleneck改成SEBottleneck。这样做的好处是基本不影响原包结构坏处是升级Ultralytics版本时需要重新添加映射。我一般会在自己项目里写一个register.py专门集中处理这些导入和映射方便统一改。3.3 直接修改源码的利与弊第三种方式最直接但也最需要储备。比如你想自定义损失函数只改yaml和新增子类都没用因为损失函数是在ultralytics/utils/loss.py里被硬编码调用的。这时候你只有两条路一是fork官方仓库在loss.py里改二是通过monkey-patch在训练脚本里覆盖原有类。我在早期项目里试过直接改源码确实方便当时想加一个辅助分类头直接在DetectionModel.forward里加了个分支很快验证了想法。但发版后遇到bug想同步上游更新时git冲突多到怀疑人生。后来我学乖了如果能用子类继承、注册覆盖的方式就不要改原文件。如果实在要改就把改动的部分抽成独立函数在源码里只留一行调用这样diff最小。4. 实战一自研轻量Backbone替换YOLOv8默认Backbone4.1 目标场景与设计思路嵌入式设备上跑缺陷检测输入图像不大但对延迟敏感。默认YOLOv8s的Backbone包含多次下采样和大量C2f对算力要求还是偏高。我的设计目标是做一个三层下采样的轻量Backbone保持特征提取能力的同时减少通道数。结构上我选择了“卷积倒残差轻量注意力”的路线。虽然Ultralytics自带了一些经典Backbone但为了让读者明白完整的替换过程我以一个自定义LightBackbone为例避免你只会改官方已有模块、遇到新结构却不会注入。4.2 编写自定义模块并注册新建light_backbone.py定义整个Backbone结构。这里不用太复杂突出“能被yaml直接实例化”的关键点import torch import torch.nn as nn from ultralytics.nn.modules import Conv class LightStem(nn.Module): 下采样stem使用两个conv快速降低分辨率 def __init__(self, c1, c2): super().__init__() self.conv1 Conv(c1, c2, k3, s2) self.conv2 Conv(c2, c2, k3, s1) def forward(self, x): return self.conv2(self.conv1(x)) class LightBlock(nn.Module): 轻量特征提取块通道减半使用深度可分离卷积 def __init__(self, c1, c2, n1, shortcutTrue): super().__init__() cv1 Conv(c1, c2, 1, 1) cv2 Conv(c2, c2, 3, 1, gc2) # depthwise cv3 Conv(c2, c2, 1, 1) self.block nn.Sequential(cv1, cv2, cv3) self.add shortcut and c1 c2 def forward(self, x): return x self.block(x) if self.add else self.block(x) class LightBackbone(nn.Module): def __init__(self, c13, c264): super().__init__() self.stem LightStem(c1, c2) self.stage1 LightBlock(c2, c2, n1) self.down2 Conv(c2, c2 * 2, 3, 2) self.stage2 LightBlock(c2 * 2, c2 * 2, n2) self.down3 Conv(c2 * 2, c2 * 4, 3, 2) self.stage3 LightBlock(c2 * 4, c2 * 4, n2) self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(c2 * 4, c2 * 4) # 仅用于输出特征实际可去掉 def forward(self, x): x self.stem(x) x self.stage1(x) x self.down2(x) x self.stage2(x) x self.down3(x) x self.stage3(x) return x注意这个自定义Backbone需要输出YOLO检测头需要的特征图列表。为了简化这里只演示注册方法实际替换Backbone时你要让不同尺度的特征图都能从Backbone中返回。我的做法是在LightBackbone.forward里返回[x_stage1, x_stage2, x_stage3]三个特征图分别对应80×80、40×40、20×20的检测尺度。然后把它注册到Ultralytics中。在tasks.py头部加入from light_backbone import LightBackbone再在parse_model可用的模块字典中补上m eval_name if m in {...}: ... else: m globals()[m]最省事的注册方式是直接把类导入到ultralytics.nn.modules.__init__.py的__all__里让tasks.py的from ultralytics.nn.modules import *自动拿到类名。我在实际项目中更愿意维护一个patching.py在训练脚本启动时执行tasks.parse_model里的字典更新避免改动包本体。4.3 修改模型yaml并验证结构现在写一个yolov8-light.yaml只保留YOLOv8的检测头替换掉整段Backbone# Ultralytics YOLOv8-light.yaml backbone: - [ -1, 1, LightBackbone, [3, 64] ] head: - [ -1, 1, Conv, [ 128, 3, 2 ] ] - [ -1, 1, C2f, [ 128, True ] ] - [ -1, 1, SPPF, [ 128, 5 ] ] - [ -1, 1, Conv, [ 64, 1, 1 ] ] - [ -1, 1, nn.Upsample, [ None, 2, nearest ] ] - [ [ -1, 3 ], 1, Concat, [ 1 ] ] ...这里做了简化真正的head需要对应Backbone输出的特征图索引。我把Backbone整体当作v8原结构中第1到第9层的替代检测头里的Concat层引用了Backbone内部输出因此在yaml中要通过多个层号串联。这个最容易错我的建议是先用一个自己写的build_model.py打印结构不要急着训练from ultralytics import YOLO model YOLO(yolov8-light.yaml) model.info()运行后如果能看到每一层的输出shape并且没有KeyError说明模块注册和yaml结构基本正确。如果model.info()报RuntimeError通常是你自定义模块的forward返回的特征图数量、通道与检测头不匹配。4.4 训练与性能对比结构验证通过后用小型数据集先跑10个epoch我一般会在train命令里加上--workers 4 --batch 16 --epochs 10观察loss是否能下降。对比指标时除了mAP50、mAP50-95还要看参数量、GFLOPs和单帧耗时。我在那个缺陷项目里换了LightBackbone后参数量降到原来的38%帧率从8帧提升到25帧代价是mAP50下降了1.2%。如果你的任务对准确率更敏感就不要盲目追求轻量化可以尝试在Block里加SE或注意力再对比。5. 实战二替换损失函数解决小目标和样本不均衡问题5.1 YOLO默认损失函数的结构很多人在自定义模块时只关注网络结构忽略了损失函数才是影响训练结果的关键。YOLOv8的默认损失在ultralytics/utils/loss.py中核心是v8DetectionLoss它由三部分组成BoxLoss边界框损失通常是CIoU加上DFLDistribution Focal Loss。ClsLoss分类损失默认二值交叉熵BCE。DFL分布焦点损失用于回归框的分布表示。如果数据集存在严重的类别不均衡或者大量小目标默认BCE分类损失往往不够用。这时候就需要替换损失函数比如换成Focal Loss。5.2 自定义FocalLoss并替换在loss.py中v8DetectionLoss的__init__里会执行self.bce nn.BCEWithLogitsLoss(reductionnone)如果你想替换成Focal Loss最简单的做法是写一个FocalLoss类然后通过继承原类覆盖bce属性import torch import torch.nn as nn class FocalLoss(nn.Module): def __init__(self, gamma1.5, alpha0.25): super().__init__() self.gamma gamma self.alpha alpha def forward(self, pred, target): pred_sigmoid pred.sigmoid() pt (1 - pred_sigmoid) * target pred_sigmoid * (1 - target) focal_weight (1 - pt).pow(self.gamma) if self.alpha is not None: alpha_t self.alpha * target (1 - self.alpha) * (1 - target) focal_weight focal_weight * alpha_t loss nn.functional.binary_cross_entropy_with_logits( pred, target, reductionnone ) return focal_weight * loss然后在训练脚本里继承官方损失类并替换from ultralytics.utils.loss import v8DetectionLoss class CustomV8Loss(v8DetectionLoss): def __init__(self, model): super().__init__(model) self.bce FocalLoss(gamma1.5)最后通过custom训练入口使用这个损失。需要注意v8DetectionLoss的__init__可能已经在后续版本中调整了bce的作用范围比如用在分类损失处也可能用于DFL因此要看清代码。更稳妥的替换点是把v8DetectionLoss里的self.bce换成self.bce FocalLoss(...)但不能影响DFL内部的BCE使用因为DFL的数学形式和分类不完全一样。我在v8.2版本上实际验证过直接替换self.bce后DFL也变成Focal形式训练loss曲线变得很奇怪。后来我单独保留了一个self.bce_cls给分类分支才解决了问题。所以替换损失函数时一定要理清代码结构不能一刀切。5.3 训练中观察loss曲线与指标变化替换损失函数后不要只看mAP。先用plotsTrue训练观察train/cls_loss、train/box_loss、train/dfl_loss三条曲线。如果分类loss明显下降变缓说明Focal Loss的gamma偏大。如果目标框回归出现问题则可能是误替换了DFL部分。我会在验证集上额外统计各类别的AP而不是只看平均mAP这样才能确认是否真的解决了小目标类别的不均衡问题。6. 替换Backbone/Head必踩的坑从报错信息到修复全链路6.1 常见报错与根因对照我整理了替换模块时最常见的报错按出现频率排序报错信息根因解决方案KeyError: MyModule自定义类没有注册到tasks.py的解析环境导入模块并在模块字典中添加名字RuntimeError: size mismatch特征图通道数或层索引不匹配打印每一层输出shape检查yaml中的层索引引用TypeError: __init__() got an unexpected keyword argumentyaml中参数列表顺序或名称与构造函数不一致对照构造函数签名调整yaml参数CUDA out of memory替换后的模块显存占用更大降低batch size用--device cpu先验证结构再上GPUAttributeError: NoneType object has no attribute shape自定义模块在某个分支上返回了None检查forward路径避免条件分支缺失6.2 一次维度不匹配的完整排查记录有一次我把新设计的Neck层接进YOLOv8-seg训练到第一个step就报size mismatch。我没急着看堆栈而是先把模型切到CPU用随机张量跑了一次前向import torch model YOLO(my-seg.yaml).model x torch.randn(1, 3, 640, 640) with torch.no_grad(): y model(x)结果提示torch.cat时形状不一致。我就在tasks.py的parse_model后面临时加了打印把每一层输出的shape打出来。最后发现是yaml里Concat层引用的是第3层输出而第3层经过我替换的Block后通道数是128预期是192。问题不是模块本身而是我在yaml的层索引写错了。修正索引后一切正常。这个经历说明自定义模块后第一件事永远是“打印结构、单步前向”而不是直接扔给训练。6.3 设备与精度相关的坑AMD显卡、半精度等Ultralytics官方现在对AMD显卡的支持主要靠ROCm和DirectML。我在AMD RX6700上跑自定义模块时遇到过torch.cat算子在ROCm后端触发异常换成CPU却没有问题。这类算子兼容性坑很隐蔽建议先用小batch在CPU上验证一遍网络前向再换AMD显卡。另外开启AMP混合精度训练后自定义模块里的某些激活函数或自定义op可能不在autocast支持列表中导致loss变成NaN。最有效的排查方法是关闭AMP也就是训练时加--amp False如果loss恢复正常那就是精度不兼容。处理方式是在自定义模块的关键计算前后手动加torch.cuda.amp.autocast(enabledFalse)或者重写其中的算子。7. 进阶玩法数据集、实例分割和多目标跟踪模块的替换与适配7.1 替换数据加载与增强模块适配自己的数据集模块替换不只是网络结构的事。很多人在训练自己数据集时发现acc上不去问题常常出在数据加载环节。Ultralytics的BaseDataset和augment.py支持自定义增强器。我习惯用Albumentations库替换部分内置增强做法是继承BaseDataset在get_image_and_label里调用自己的增强管线。替换数据增强模块时要特别注意增强后需要同步更新bbox坐标比如随机旋转、随机裁剪都会改变框坐标。如果不小心只增强图像不改框模型就会学到错误映射。7.2 实例分割任务中的检测头替换YOLOv8-seg和YOLOv11-seg的检测头是Segment类它在普通Detect头基础上增加了一个掩码分支。如果你想把检测头替换成自定义分割头需要同时修改yaml的head部分和tasks.py里的Segment组装逻辑。分割头通常还要输出原型掩码系数和检测分支的边界框特征结合。很多人在替换分割头时只改了nn.Module却忘记修改v8SegmentationLoss里的掩码损失计算方式结果训练时mask_loss恒为0。我在项目里通常保留原损失函数先替换主干网络等效果稳定再动分割头。7.3 多目标跟踪模块的接入与替换Ultralytics的预测接口里集成了ByteTrack、BoT-SORT等跟踪器。如果你有自研的跟踪器可以在ultralytics/trackers/track.py中替换on_predict_start中的tracker实例。常见做法是继承BaseTracker在update方法中接受检测结果和嵌入特征返回跟踪ID。替换跟踪模块时最难处理的是“检测框和ReID特征如何对齐”因为YOLO的推理结果可能包含多个batch。我的经验是先关闭batch即predict(source..., batch1)能省掉大量调试时间。8. 版本升级时模块变动的影响与兼容性处理8.1 v8到v11哪些模块API变了Ultralytics版本迭代很快从v8到v11yaml里的模块名和代码结构都有变化。比如v11引入了C3k2、C2PSA部分旧模块虽然保留但不再是默认推荐。tasks.py解析逻辑也在变化之前使用的model.model[-1]等索引方式可能会失效。如果你用的是几个月前的自定义模块升级后很可能出现两类问题一是模块名被移除二是构造函数参数比如增加了一个c3参数导致yaml参数列表不兼容。我的建议是升级前先看官方仓库的Release Notes重点看nn/modules.py和loss.py的diff。不要直接升级大版本先在虚拟环境里跑一次model.info()和单step训练确认自定义模块的状态。8.2 自定义模块跨版本迁移方法把自定义模块从源码中抽离出来是降低升级痛苦最有效的手段。我在自己的项目里维护了一个extension/目录里面每个模块都是独立的nn.Module并且提供一个register()函数负责在Ultralytics的模块映射表中插入自定义类名。迁移时只需要做三件事确认新版本的目标类构造函数签名是否变化。修改register()里的映射关系。运行全部训练和推理测试。只要你没有改动官方源码升级时直接从Git拉取新版本即可。如果确实需要改官方源码我强烈建议用git format-patch保存一个独立补丁升级时git apply后再清理冲突而不是直接把所有改动提交到官方仓库分支里。我在多个项目里用过Ultralytics YOLO之后最大的体会是模块自定义与替换并不难难的是搞清楚它背后的加载机制和参数约定。只要你愿意花半天时间读完tasks.py里的parse_model函数再动手写第一个自定义模块后面所有替换工作都会变得非常顺手。另外自定义模块一定要写在独立文件里而不是直接改modules.py否则版本升级时会特别痛苦。如果你正准备给YOLO换Backbone或调损失函数不妨从上面的一个小案例开始先跑通再深入。
