★ FEATURED ARTICLE 前段时间一直在啃YOLO源码尤其盯着ultralytics.models.yolo.detect这个模块看了很久。很多人训练、验证、推理都直接调model.train()、model.val()、model.predict()但一旦涉及自定义loss、改数据加载逻辑、或者想精确控制推理后处理就不得不到这三段代码里翻找。这篇文章就围绕train.py、val.py、predict.py三个子模块把它们各自的核心类、方法调用链和设计逻辑完整过一遍顺便把我踩过的坑一起列出来。1. 先搞清楚detect模块的整体定位1.1 从import到运行三个类各管一段在ultralytics/models/yolo/detect/__init__.py里只做了三件事把训练、验证、预测三个类暴露出来。DetectionTrainer、DetectionValidator、DetectionPredictor分别对应YOLO检测任务的三条主流程。用下面的等价写法就能直观感受到它们的关系from ultralytics.models.yolo.detect import DetectionTrainer, DetectionValidator, DetectionPredictor # 等价于 from ultralytics.models.yolo.detect.train import DetectionTrainer from ultralytics.models.yolo.detect.val import DetectionValidator from ultralytics.models.yolo.detect.predict import DetectionPredictor从框架整体看ultralytics包的所有任务检测、分割、分类、姿态共用一套BaseTrainer、BaseValidator、BasePredictor骨架检测任务只是把关键方法重写掉。所以读detect模块的代码本质上是在看“检测任务如何实现Base类里定义的抽象流程”。这三条流程的职责可以用一张表概括类继承来源核心职责主要入口方法DetectionTrainerBaseTrainer构建模型、加载数据、执行训练循环、保存checkpointtrain()DetectionValidatorBaseValidator构建验证数据、前向推理、计算指标、输出混淆矩阵val()DetectionPredictorBasePredictor加载模型权重、预处理输入、推理、后处理解码、结果可视化predict_cli() / stream_inference()读代码时我建议从predict.py入手它的方法最少、链路最短读完再回头读train.py会顺畅很多。1.2 TaskType和模型工厂的映射关系detect模块虽然独立但它通过模型工厂暴露给了用户。ultralytics/models/__init__.py里根据task参数决定返回哪个TaskTypeclass TaskType(Enum): DETECT detect SEGMENT segment CLASSIFY classify POSE poseDetectionTrainer等类不会直接在用户代码里被实例化而是由model.train()内部根据taskdetect自动路由过去。训练时你传入的yolo11n.yaml等配置会同时决定模型结构和trainer类型这条链路的起点在ultralytics/engine/model.py。实际调试时可以手动实例化trainer绕过model.train()高层封装这样能直接看到每个环节做了什么trainer DetectionTrainer(overrides{model: yolo11n.yaml, data: coco8.yaml}) trainer.train()这里有个值得注意的设计所有配置参数如epochs、batch、imgsz都通过overrides字典传入由cfg机制统一管理。这样做的好处是命令行、Python API、配置文件三套入口共用一套参数校验逻辑。2. DetectionTrainer训练核心流程拆解2.1 get_model从yaml到nn.Module的构建过程DetectionTrainer中get_model负责把模型配置文件变成可训练的nn.Module。YOLO11的模型构建走的是“yaml解析 动态生成”路线def get_model(self, cfgNone, weightsNone, verboseTrue): model DetectionModel(cfg, ch3, ncself.data[nc], verboseverbose) if weights: model.load(weights) return modelDetectionModel在解析yaml时会把配置里的每一层结构先存成字典列表再通过parse_model递归生成逐层网络结构。这部分代码初读时容易懵因为它大量使用了globals()动态查找层类# 简化示意 for i, (f, n, m, args) in enumerate(d[backbone] d[head]): m eval(m) if isinstance(m, str) else m # ... layers.append(nn.Sequential(*[m(*args) for _ in range(n)]))这种动态构建方式的好处是新增一个层只需要在nn.modules里定义类然后在yaml里按配置写即可模型结构完全和代码解耦。代价是IDE跳转时会比较痛苦建议配合# type: ignore和pdb去看每一层的实际类型。nc类别数会在模型配置和数据集配置之间做校验。如果yaml里写了nc: 80而data配置的类别数是别的值DetectionModel会直接用data的nc覆盖模型yaml里的定义。这就是为什么你改了数据集的类别数模型输出层会自动调整的原因。提示自定义模型结构时detect头部的cv2、cv3分支数量必须和nc匹配。如果你看到shape mismatch错误十有八九是模型yaml和数据集类别数没对齐。2.2 set_model_attributes与加载预训练权重set_model_attributes这个方法很朴素核心逻辑是把数据集信息绑定到模型上def set_model_attributes(self): self.model.nc self.data[nc] self.model.names self.data[names] self.model.args self.args这里的names是类别名列表会在推理时用来渲染标签。args是训练超参数会保存到checkpoint里方便恢复训练时保持一致性。预训练权重加载的逻辑也很直接if weights and weights.endswith(.pt): ckpt torch.load(weights, map_locationcpu) model.load_state_dict(ckpt[model].float().state_dict())实际踩坑记录如果你用model DetectionModel(cfg, ch3, nc80)后再调用load()加载COCO预训练权重但自定义数据集只有两个类别比如person和car那么加载权重会直接报shape mismatch。解决办法是用官方的建议设置transferTrue它会自动跳过shape不匹配的层只加载能对齐的部分。2.3 get_dataloader和build_dataset训练数据流水线build_dataset看名字简单内部做了不少事。检测任务训练时的数据集构建主要有三个分支非矩形推理rectFalse普通YOLODataset每张图独立缩放后拼batch矩形推理rectTrue同一个batch里的图片宽高比尽量一致减少padding浪费是否启用mosaic增强启用时每4张图拼成1张再随机裁剪相当于用显存换数据多样性具体代码里关于mosaic的判断值得注意if self.args.rect and mf: # 矩形推理模式下若mosaic系数为0 dataset YOLODataset( img_path... augmentFalse, rectself.args.rect, ... ) else: dataset YOLODataset( img_path... augmentTrue, rectFalse, ... )读代码时我这里的理解是矩形训练模式下每张图的输入尺寸不固定而mosaic增强需要先把多张图拼成固定尺寸再裁剪两者逻辑冲突。所以框架在启动矩形训练时会强制关闭mosaic这是很多人在自定义数据流程时忽略的前提条件。如果你中途打开rectTrue但没注意这个分支训练出的模型在验证时bbox偏移会很明显。get_dataloader比较简单就是包了一层DataLoaderreturn self.dataset if self.args.v5loader else DataLoader(...)那行v5loader是历史遗留老项目沿用v5的数据加载方式新版默认走YOLODataset自己的__getitem__逻辑。2.4 训练主循环loss反向传播与老代码的差异detect模块的train.py不直接写训练循环继承自BaseTrainer。BaseTrainer里已经有一个完整的train方法包含for epoch in range(start_epoch, epochs)循环、学习率调度、EMA更新、梯度裁剪、checkpoint保存等。DetectionTrainer主要在get_model、get_dataloader、build_dataset这些模板方法上和检测任务对齐loss计算在DetectionModel的forward内部完成。新版YOLO的loss计算已经内嵌到模型内部不是传统的“模型输出y_pred再在trainer里算loss”。模型forward返回一个Loss对象内部包含分类loss和回归lossloss, loss_items model(batch)loss_items是独立loss明细会被记录到日志中训练时你可以观察到box_loss、cls_loss、dfl_loss各自的变化趋势。这种设计的直接好处是验证阶段模型前向输出的是preds训练阶段输出的是loss同一个forward入口通过training标志区分。我调试时常用的一个方法是手动给loss加断点看分类分支和回归分支分配到的梯度比例。如果是类别极度不平衡的数据集比如只有一个大类“cls_loss”会一直降不下去这时就需要调整cls和box在loss中的权重系数粮食都在模型的__init__.py附近。3. DetectionValidator验证环节实现解读3.1 验证集的数据构建与训练集的差异val.py里的build_dataset和get_dataloader与训练版本相比有本质区别——验证集不需要数据增强这是铁律加了增强指标就废了。验证数据的关键是必须保证每张图都能被正确处理填充方式要和训练保持一致。def build_dataset(self, img_path, modeval, batchNone): gs max(int(de_parallel(self.model).stride.max()), 32) return YOLODataset( img_pathimg_path, imgszself.args.imgsz, batch_sizebatch, augmentFalse, rectself.args.rect, stridegs, pad0.5, prefixf{mode}: , ... )这里stridegs的作用是把图片padding到模型最大stride的整数倍。比如模型最大stride是32输入尺寸为640图片会被padding到320或640的整数倍。这个细节直接关系到你验证时输入分辨率是否真的等于设置值。我遇到过自己预处理时没对齐stride导致推理结果出现边框错位的情况。3.2 模型前向与后处理解码DetectionValidator里的postprocess把模型输出的原始prediction张量转换成可以评估的bbox格式。def postprocess(self, preds): return ops.non_max_suppression( preds, self.args.conf, self.args.iou, labelsself.lb, multi_labelTrue, agnosticself.args.single_cls, max_detself.args.max_det, )non_max_suppression是推理链路的核心瓶颈之一它内部会按置信度排序计算IoU递归消除重叠框。验证阶段用labelsself.lb传入GT框参与NMS计算目的是让预测框尽量贴近GT框评估指标会更接近真实部署场景。推理阶段则不会传labels。后面还会看到一行关键调用preds self.model(batch[img])验证代码里没有显式先model.eval()因为BaseValidator已经处理了。但如果你自己写评估脚本记得手动执行model.eval()和torch.no_grad()否则BN层的统计量会在验证集上被错误更新指标看起来会很奇怪。3.3 update_metrics与指标计算的隐藏逻辑验证的核心指标计算全部在update_metrics方法中。这个方法接收模型输出preds和GT batch通过匹配预测框和真实框来更新混淆矩阵、计算TPTrue Positive、FPFalse Positive、FNFalse Negative。这些状态存在self.metrics和self.jdict里。def update_metrics(self, preds, batch): # 实现NMS后处理与GT配对 # 更新混淆矩阵、AP/Target等统计 ...指标类型上YOLO的验证会同时维护两类self.metrics最终mAP结果每张图预测结束后累计self.jdict每个类别的AP统计明细验证结束时finalize_metrics会把累计数据做平均计算得到最终mAP50、mAP50-95等数值。看源码时一定记住一个坑update_metrics中使用的是pbatch即处理后的batch其中“batch的img是预处理后归一化、填充后的张量”真正评估处理时会把bbox的坐标缩放到原图尺寸所以验证精度不受预处理影响。如果你准备自定义评估函数这个坐标系变换是绕不开的。4. DetectionPredictor推理链路解析4.1 preprocess从一张图到模型输入张量predict.py的核心是DetectionPredictor类继承自BasePredictor。推理前最重要的方法就是preprocess它把输入图像转换成模型需要的张量格式def preprocess(self, img): # 如果没有resize则保持原图尺寸但做letterbox填充 if not isinstance(img, torch.Tensor): # 输入是路径、numpy数组或PIL图像 img torch.from_numpy(img).to(self.model.device) img img.float() / 255.0 # 归一化到0~1 # 调整维度顺序BCHW并letterbox处理 return img这里我插一句很多人以为YOLO推理时会resize实际是letterbox等比缩放加灰边。preprocess中如果图片尺寸不是模型输入尺寸的整数倍会做边缘padding。这个padding在你自己实现部署推断逻辑时极易漏掉chartbox的偏移误差就直接出现了。实际代码中letterbox的实现通常会把变换矩阵保存下来用于后处理时把bbox坐标映射回原图。所以preprocess不只是简单resize它输出了一个可用于反变换的scale、pad信息这是许多自研推理代码会忽视的地方。模型的输出直接就是NMS之前的原始prediction。所以你在用model.predict()的时候其实框架已经内部完成了候选框解码、NMS、映射回原图坐标这一整套。但如果你在自己写推理循环需要手动调用model.forward()拿原始输出然后自己做non_max_suppression。4.2 postprocess与结果展示DetectionPredictor的postprocess主体逻辑是NMSdef postprocess(self, preds, img, orig_imgs): preds ops.non_max_suppression( preds, self.args.conf, self.args.iou, agnosticself.args.agnostic_nms, max_detself.args.max_det, classesself.args.classes, ) # 再把坐标映射回原始图像尺寸 ... return predsorig_imgs参数是原始图像或者原始图像的尺寸信息保留下来用于还原坐标。推理结果框的坐标坐标系发生了什么转化需要理清楚模型输入坐标是在letterbox后的图上归一化的输出是letterbox图上坐标映射回原始图尺寸靠postprocess里的反变换实现。调用链上BasePredictor.stream_inference是推理的实际执行函数它循环遍历数据源对每个batch图像调用inference模型前向和postprocess后处理。如果开了streamTrue推理时source可以是视频流、摄像头或迭代器这样可以边读帧边处理避免视频推理卡顿。4.3 推理结果返回格式解析现代版YOLO的统一推理入口是model.predict()它返回一个Results对象列表。每个Results对象内部包含了boxes检测框xyxy格式、置信度、类别索引masks分割掩码如果是分割模型keypoints关键点如果是姿态模型orig_img原始图像数组names预测类别对应的名称你可能会问如果直接看源码postprocess里返回的是元组(preds, boxes, masks)为什么用户看到的是Results对象关键在于BasePredictor在完事后调用了self.results self._postprocess_after_inference之类的逻辑把原始预测打包成了用户友好的Results对象。这部分在ultralytics/engine/predictor.py里。5. 三个模块的联动与我最常遇到的坑5.1 从Base到子类的调用链全览虽然三个Python文件分别实现了train、val、predict但它们在运行时是互相耦合的。训练开始前DetectionTrainer会做校验并构建DetectionValidator训练中途每隔一段时间会自动跑一次验证保存mAP最高的权重推理加载模型时又会根据模型类型自动选择对应的DetectionPredictor。# 训练中创建验证器的核心逻辑 validator self.get_validator() validator(modelself.model, dataloaderself.get_dataloader(...), save_dirself.save_dir)所以你在train.py里会看到get_mot_validator之类的方法名本质上就是返回DetectionValidator。在val.py里也有get_model等手脚。这是经典的Template Method模式Base类定义算法骨架子类实现具体步骤。检测、分割、分类各任务互不干扰新增任务只需要继承Base并实现对应方法即可。调试时给个实用建议如果在训练过程中想单独验证某个checkpoint可以这样写model YOLO(runs/detect/train/weights/best.pt) model.val(datacoco8.yaml, splitval, imgsz640)这会自动走DetectionValidator流程和你训练中途的验证完全一致。5.2 高频报错与排查心得先说一个我调试YOLO时踩得最深的坑DataLoaderworker数量。如果你在get_dataloader里看到workersself.args.workers而训练机是Windows系统workers设置过大时比如8经常出现假死或报BrokenPipeError。通常我建议Windows用户直接设workers2或0Linux可以按CPU核心数灵活设但要避开机器IO瓶颈。第二个坑在验证阶段val.py计算mAP时会要求GT数据格式严格对齐否则虽然不报错但mAP值会低得离谱。常见问题是data配置里的names顺序和训练时不一致。YOLO的GT坐标格式是归一化xywh中心点宽高输入图像尺寸、letterbox padding对坐标的影响已经体现在前面说的pad参数中。如果觉得验证框和GT对不上优先检查这个。第三个坑是predict.py里conf参数。默认conf0.25拦掉了不少低置信度的框如果你用同一个模型做高召回要求的业务比如安全帽检测单纯提高iou没用要降低conf到0.05甚至0.01。代价是NMS计算量上升因为待选框数量变多了。这里我实测下来conf0.1、iou0.45在大多数场景下是召回和精度的良好平衡点。问题现象大概率原因排查方法训练时loss为NaN学习率过大 / 数据有异常标注比如坐标超出图像范围查看cfg里的lr参数打印batch里的归一化坐标检查是否在[0,1]区间推理输出空列表conf阈值设置过高 / 模型未正确加载权重将conf降到0.01打印原始preds看是否存在高置信度框验证mAP比训练mAP低很多验证数据增强开启 / 图像尺寸设置不一致检查val.py里augment是否为False是否imgsz和训练时相同内存溢出batch size过大 / 启用了mosaic但数据集中有超大图调小batch、关闭mosaicmosaic0或将图片预缩放到较小尺寸5.3 自定义修改的建议路由很多初学者会直接改train.py里的大段代码我建议不要这样操作。Ultralytics的设计本来就支持通过继承和hook机制来做自定义修改想要改数据增强逻辑继承YOLODataset重写__getitem__想要改loss继承DetectionModel在forward里替换loss模块想要改训练循环的某个节点继承DetectionTrainer重写对应hook方法之前我改造某个两阶段检测任务时就是直接继承DetectionTrainer重写get_dataloader和build_dataset。这种做法的好处是不需要动官方包一行业代码未来升级Ultralytics版本时直接无缝衔接。在train.py的get_model中还有一个细节如果使用cfg模式yaml创建模型nc参数会直接用数据集类别数覆盖模型配置文件如果直接传入weights模式.pt则会先构建模型再覆盖权重层两种方式的加载顺序是不同的。这个差异在迁移学习场景下体现得很明显——直接传.pt作为模型时nc会以模型文件里的为准除非你显式设置data里的nc。val.py里get_model通常都是空方法因为验证时模型已经从checkpoint里恢复了不需要重建。但分割模型有所不同它需要在get_model里指定mask分支的通道数。predict.py还有一个容易被忽略的stream_inference方法。这是推理的主循环负责迭代source可以是图片路径、视频文件、摄像头、PIL图像、numpy数组等对每个对象做预处理、推理、后处理、结果记录。如果你要实现自定义推理服务大部分情况重写preprocess和postprocess就够了主循环逻辑用官方实现。6. 三联表三个文件的设计对比与性能提示最后用一张表总结三个文件的设计差异列表在快速阅读时真的有用对比维度train.pyval.pypredict.py核心类DetectionTrainerDetectionValidatorDetectionPredictor关键hookget_model / build_dataset / get_dataloader / set_model_attributesbuild_dataset / get_dataloader / update_metrics / finalize_metricspreprocess / inference / postprocess输入数据配置 模型配置 超参数checkpoint 数据配置 验证参数任意图像源 模型权重输出权重文件 训练日志 EMA模型mAP指标 混淆矩阵 预测可视化Results对象列表或保存的图片/视频支持自定义的侧重点数据流 / 网络结构 / 训练超参指标计算方法 / 匹配逻辑 / 数据加载输入格式 / 后处理策略常见性能瓶颈mosaic数据增强 多worker加载NMS和GT匹配大验证集时NMS letterbox预处理读这些代码时我一直觉得最有意思的设计是让训练、验证、推理共用同一个DetectionModel而不同阶段通过传入的training标志或环境变量来切换行为。模型内部会build出不同的分支loss只在训练时计算验证时只输出预测张量。这让你在做模型转换或量化感知训练时可以只关注模型结构而不用操心流程层次的切换。最后分享一个我调试时的习惯每拿到一个YOLO版本我会先跑一遍三行代码把所有hook的调用顺序和输入输出打印出来。from ultralytics import YOLO model YOLO(yolo11n.pt) print(model.predict(sourcebus.jpg, conf0.1, verboseTrue)) model.train(datacoco8.yaml, epochs1, imgsz64, workers0, verboseTrue) model.val(datacoco8.yaml, imgsz64, verboseTrue)用verboseTrue打开日志输出观察训练流程的打印顺序能快速定位自定义修改应该插在哪一步。这种“先跑通-再打断点-再改造”的路线比一上来就啃每行代码的效率高很多特别是对于初学者来说强烈建议照搬这个习惯。
阅读完成 · 觉得有帮助？
