简介这是一份YOLOv8目标检测自定义数据集训练源码工程面向有一定深度学习基础、需将模型迁移到自身业务场景的开发者与研究人员覆盖工业质检、安防监控、遥感识别等典型落地场景。压缩包共24个文件主体为23个Python脚本并附1个txt说明文档整体仅51KB虽小巧但训练链路完整。脚本从数据准备入手包含数据集制作与检查、数据加载、标签分配再到训练、验证、推理与评估的一整套流程同时集成CBAM注意力、损失重加权及Transformer结构等改进组件既可直接训练也可作为算法调优的起点。工程目录按模块拆分便于按图索骥修改网络结构或损失函数读者替换为自己的图像与标注后即可启动训练也能借源码梳理YOLOv8的训练链路设计。目前已有4495人学习对希望快速上手自定义训练或研究检测改进点的读者具有较高参考价值。1. yolov8目标检测训练自己的数据集这份源码比官方demo多做了哪几件事不少朋友跑通官方YOLOv8 demo之后以为“训练自己的数据集”只是把图片路径换一下的事结果一上手就卡在 loss 不降、mAP 上不去、训练到一半显存爆掉这类问题上。这份 yolov8 目标检测训练自己的数据集源码解决的问题不是“怎么跑通 demo”而是“怎么把一次训练做成可复现、可回滚、可查错的完整工程流程”。源码作者做了十年 YOLO 相关算法仿真包里除了常规的 train.py 和 predict.py还拆出了 check_dataset.py、loss_reweighting.py、CBAM.py、label_assignment.py、parse_weight.py 这些针对训练全流程的定制模块。适合正在准备自己数据集、想系统落地一次目标检测训练、以及想拿官方实现做对照二次开发的人。2. 数据准备与校验先把 make_dataset.py 和 check_dataset.py 跑通2.1 数据集目录结构与标注格式先统一再开工训练自己的数据集八成的问题出在数据格式上。YOLOv8 官方要求的目录结构是 images 和 labels 两个平级目录每张图片对应一个同名 txt 文件txt 每一行是class cx cy w h其中 cx、cy、w、h 都是相对图片宽高的归一化值。这个包里的 make_dataset.py核心就是把 VOC 的 xml 或 COCO 的 json 转成这种 txt 格式。我一般会先把原始数据分成三类目录JPEGImages 放原图Annotations 放标注ImageSets 放划分好的 train.txt / val.txt 文件清单。常见的公开数据集里CCPD 车牌数据集和 HRSC2016 遥感舰船数据集的标注格式都不相同CCPD 的文件夹名里直接带 bbox 信息HRSC2016 用的是 xml 加额外旋转框字段。不统一转换就直接训练第一个报错往往就是 “Label class ID out of range” 或 “all labels are empty”。推荐目录结构如下my_dataset/ ├── images/ │ ├── train/ │ └── val/ ├── labels/ │ ├── train/ │ └── val/ ├── train.txt ├── val.txt └── dataset.yamldataset.yaml 里只需要写三类信息train 和 val 的图片路径、类别数量 nc、类别名字 names。这份源码在读取时依赖这个 yaml 文件定位数据集所以路径建议写绝对路径避免不同机器上相对路径解析不一致。2.2 make_dataset.py 怎么用从 VOC 标注到 YOLO 格式从文件名看make_dataset.py 就是干这个转换的。我看了源码结构它至少做了三件事读取标注文件xml/json/txt、把 bbox 坐标归一化、按比例划分 train/val。命令行参数大致是这样python make_dataset.py \ --source ./JPEGImages \ --annotations ./Annotations \ --output ./my_dataset \ --split 0.8 \ --classes plane car ship参数说明--source是原始图片目录--annotations是标注目录--output是转换后数据集要放的位置--split是训练集占比0.8 意味着 80% 做训练、20% 做验证--classes是类别名列表顺序就是类别 ID 的映射顺序。这四个参数里最容易出问题的是--classes的顺序——类别 ID 由这个参数的传参顺序决定不是由标注文件里的字符串自然排序决定。如果原始标注是 VOC 格式转换的核心逻辑就是读 xml 里的bndbox字段然后归一化import xml.etree.ElementTree as ET def voc_to_yolo(xml_path, width, height, class_map): tree ET.parse(xml_path) root tree.getroot() lines [] for obj in root.findall(object): name obj.find(name).text if name not in class_map: continue cls_id class_map[name] box obj.find(bndbox) x1 float(box.find(xmin).text) y1 float(box.find(ymin).text) x2 float(box.find(xmax).text) y2 float(box.find(ymax).text) w (x2 - x1) / width h (y2 - y1) / height cx ((x1 x2) / 2) / width cy ((y1 y2) / 2) / height lines.append(f{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}) return lines这段代码的逻辑是遍历 xml 里每个 object查类别名在 class_map 中的 ID然后算归一化后的中心点坐标和宽高。注意两个细节宽高归一化要用原图的宽高不能用标注的宽高cx和cy是中心点不是左上角很多新手在这里少除一个 2。转换完以后建议抽查几张图把 txt 里的坐标画回图上确认框的位置没有偏移这一步能省掉后面无数排查时间。2.3 check_dataset.py训练前最后一道体检这个文件在包里单独存在我认为是整份源码里性价比最高的一个模块。官方 ultralytics 仓库里也有一段检测数据集合法性的逻辑但那个跑完只告诉你数据集能不能用不会告诉你哪里有问题。这个 check_dataset.py 更偏“自己训练的体检”我拆出来看它至少检查四类东西图片能不能正常打开、标签文件有没有空标签、坐标有没有越界小于 0 或大于 1、类别 ID 有没有超出 nc 范围。跑完它之后你会拿到一份统计报告大概长这样检查项通过数异常数说明图片可读性124533 张图片损坏或非标准编码空标签文件12451212 张图没有对应 txt 或 txt 为空坐标越界124555 个 bbox 的 w/h 大于 1类别 ID 越界12450类别 ID 都在 nc 范围内从实际经验看空标签文件的危害被严重低估。一张图没有标签训练时它不会报错但它会成为纯背景样本如果这种样本太多模型会偏向输出大量低置信度框或者干脆把背景学成目标。坐标越界更隐蔽w 或 h 大于 1 的框在 loss 计算时不会立刻让训练崩溃但 mAP 会莫名上不去。我现在的习惯是每次换数据集先跑一遍 check_dataset.py 再决定要不要训练这份源码帮你把这一步强制前置了。3. 训练主流程train.py 和 main.py 的参数体系怎么调3.1 启动命令与核心参数这个包的入口是 main.py训练逻辑在 train.py 里。启动训练的命令大致是python main.py \ --mode train \ --data ./my_dataset/dataset.yaml \ --cfg ./modules/yolov8.yaml \ --epochs 100 \ --batch 16 \ --img-size 640 \ --device 0参数含义拆开讲--mode控制是训练还是推理train 和 predict 两个值--data指向刚才说的 dataset.yaml--cfg是模型结构配置这个包把模型结构拆到了 modules 目录下官方实现的配置文件也兼容--epochs和--batch不多说--img-size是训练输入尺寸640 是默认值显存小可以降到 512但会略微影响小目标检测效果--device是 GPU 编号。环境配置这里有一个容易被忽略的点。requirements.txt 里的依赖版本是作者在他自己机器上验证过的不要盲目升级到最新版尤其是 torch 和 torchvision 的版本组合。我见过两次翻车案例第一次是 torch 升到 2.x 后某个算子在新版本里改了默认行为loss 曲线直接变成一条直线第二次是 CUDA 版本和 torch 不匹配训练时 CPU 跑满但 GPU 利用率只有个位数。你的 torch 版本和机器 CUDA 版本要对得上训练前跑一段torch.cuda.is_available()验证一下。3.2 schedule.py 学习率调度让训练过程更可控train.py 里的训练循环本身没太多定制空间真正影响收敛效果的往往在学习率策略上。schedule.py 实现的是 warmup 余弦退火的组合这是目标检测训练里最常见的做法前几个 epoch 用较小的学习率让模型稳定起步然后进入余弦退火阶段逐步降低学习率。默认参数下warmup 是 3 个 epoch初始学习率是 0.01。如果你发现 loss 在开头几个 epoch 就剧烈震荡优先怀疑初始学习率过大把 0.01 改成 0.005 再看。如果你的数据集很小几百张图warmup 可以放大到 5 个 epoch让模型先适应你的数据分布。这个文件里应该还有 schedule 曲线可视化的入口盲看数字不如画一条学习率曲线直观。余弦退火的结束学习率一般设置到初始学习率的 1/100 到 1/1000。这个值直接影响最终收敛精度设得太大后期 loss 会在一个平台上抖动设得太小最后若干 epoch 几乎白跑。我习惯先跑 100 epoch 的默认配置保存 loss 历史再根据曲线尾部形态决定要不要调。3.3 loss_reweighting.py 和 CBAM.py从损失重加权到注意力机制loss_reweighting.py 解决的是类别不均衡问题。目标检测里经常遇到一类目标几百张、另一类只有几十张的情况模型会天然偏向样本多的类别。这个文件实现的常见做法是 inverse frequency reweighting即按类别样本数倒数归一化后作为 loss 权重import numpy as np def compute_class_weights(label_files, num_classes): counts np.zeros(num_classes) for txt_path in label_files: with open(txt_path, r) as f: for line in f: cls_id int(line.split()[0]) if cls_id num_classes: counts[cls_id] 1 counts np.maximum(counts, 1.0) # 防止除零 weights 1.0 / (counts / counts.sum()) weights weights / weights.sum() * num_classes # 归一化到均值1 return weights逻辑就是统计每个类别出现的总次数按反比算出权重最后归一化。注意倒数那一步counts 为 0 的类别如果硬除会出 inf所以先取 max 把 0 拉到 1。这个权重会在 loss 计算时乘到对应类别的分类损失和回归损失上类别越少权重越大相当于强行拉平模型对不同类别的关注度。如果你的数据集本身类别分布均匀这个模块可以不用如果明显不均衡开着比调任何超参数都管用。CBAM.py 实现的是通道注意力 空间注意力模块。YOLOv8 的 backbone 是 C2f 结构CBAM 一般插在特征提取的后半段或 neck 部分让模型更关注“哪里有什么”。这个文件是给想改结构的熟手准备的直接在 modules/yolov8.yaml 里给某个层后面加一行 CBAM 配置就能启用。要注意的是CBAM 不是加了就变强——小数据集上加注意力模块模型参数量增加但数据不够mAP 可能反而下降。4. 标签分配、权重解析与云端训练源码包的三个隐藏模块4.1 label_assignment.py正负样本是怎么分配的YOLOv8 官方默认用的是 task-aligned assigner核心逻辑是同时考虑分类得分和 IoU 来挑选正样本。这个包里把 label_assignment 单独拆成一个文件说明作者大概率改过分配逻辑或者想让这一步可以被替换。简单理解标签分配要做的事是每个 GT 框从所有预测框里选哪些作为正样本去计算 loss哪些作为负样本。常见做法是算每个预测框和 GT 的 IoU超过阈值的当正样本低于阈值的当负样本。YOLOv8 在此基础上引入了对齐度指标把分类置信度和 IoU 相乘值越大越容易被当成正样本。为什么这个文件值得看因为如果你处理的是小目标数据集默认的分配策略可能让很多小目标匹配不到正样本导致 loss 看着在降、小目标一个都检不出来。动手调这个文件时关注两个值topk 的选择范围和 IoU 下阈值。topk 越大正样本越多训练越稳定但可能引入低质量匹配下阈值越低负样本越少。我一般会在小目标数据上把 topk 从 13 调到 15然后观察小类别的 recall 变化。4.2 parse_weight.py预训练权重加载与续训的格式适配parse_weight.py 这个文件我判断它是做两件事一是把官方 yolov8.pt 预训练权重解析成这个框架能读的字典格式二是在类别数不一致时对权重做裁剪或随机初始化。当你用自己的数据集训练类别数多半不是 COCO 的 80 类。比如 CCPD 车牌检测是 1 类HRSC2016 舰船检测也是 1 类。预训练权重的最后一层输出维度是 80和你的数据集不匹配直接加载会报 shape 不匹配。parse_weight.py 的常见做法是遍历预训练权重字典跳过 head 层的参数或者只保留 backbone 和 neck 部分import torch def parse_pretrained_weight(ckpt_path, num_classes, model_state_dict): raw torch.load(ckpt_path, map_locationcpu) raw_state raw.get(model, raw) new_state {} for key, value in raw_state.items(): # 丢弃属于分类头的层这些层维度依类别数而定 if cls in key or cv3 in key: continue if key in model_state_dict and model_state_dict[key].shape value.shape: new_state[key] value # 对分类头做随机初始化 for key in model_state_dict.keys(): if cls in key or cv3 in key: new_state[key] model_state_dict[key] return new_state逻辑说明遍历预训练权重凡是名字里带 cls 或 cv3 的层分类头和检测头直接丢弃因为输出类别数量不同剩下的层如果形状一致就复用预训练值最后把分类头替换成当前模型随机初始化的参数。这个做法比直接加载再微调更稳避免分类头的错误维度影响反向传播。一个容易踩的坑是有的权重文件里 key 带module.前缀DataParallel 保存的有的不带。parse_weight 里应该做了这个 key 的清洗如果你自己写加载逻辑要记得strip_prefix一下否则所有 key 都对不上模型会用随机权重从头训练等于预训练完全白搭。4.3 moxing_from_obs.py 与 util 下的工具云端训练与评估moxing_from_obs.py 是华为云 ModelArts 场景的文件。如果你在本地训练这个文件用不上但如果你的训练数据存在 OBS 上训练前需要先把它同步到本地缓存。这个文件的逻辑一般是调用 moxing 的 copy 接口把 OBS 路径的数据拉到本地 tmp 目录再走正常的数据加载流程。util 目录下几个文件值得提。misc.py 和 common.py 应该是通用工具函数比如 IoU 计算、bbox 格式转换metrics.py 算 mAP 和 recallcoco_eval.py 是 COCO 评估协议的实现aisafety_model_eval.py 看起来是作者所在业务场景的评估脚本可能是针对特定安全类目标的指标统计。这些文件在训练完的验证阶段会被调用predict.py 输出检测框之后就是靠这些模块算出你真正关心的指标。如果你的场景是通用目标检测用 metrics.py 和 coco_eval.py 就够了aisafety 那个文件可以先不管。5. 避坑训练自己的数据集最容易翻车的五个场景5.1 loss 不降且震荡标签文件和类别 ID 对不上现象训练了十几个 epochloss 一点下降趋势都没有甚至来回震荡。打开日志看分类 loss 和回归 loss 都在一个高值附近抖动。原因八成是标签文件里的类别 ID 和数据集配置对不上。比如 make_dataset.py 里 --classes 传的是plane car ship但某个标注文件写的是1 0.5 0.5 0.2 0.2其中 1 对应的是 car而 txt 里的类别 ID 越界或映射错了。也可能是一些样本的标签文件确实是空的模型在这些图上只能学到背景。解决跑一遍前面说的 check_dataset.py它会列出所有类别 ID 超出范围和标签为空的文件。逐个修正后重新生成数据集。这个坑我刚开始训练时踩过当时排查了两天才发现是一个子目录的标注文件是从旧项目拷过来的类别定义完全不一样。5.2 mAP 卡住不动小数据集强行堆模块现象训练过程正常loss 正常下降但 mAP 到 0.4 左右就上不去了怎么调学习率都白搭。原因数据量本身不够或者模型结构加了太多额外模块。CBAM 和 Transformer 这类模块不是免费的——它们增加参数量需要更多数据来拟合。几百张图的数据集硬上这些模块模型很容易过拟合到训练集验证集 mAP 自然卡住。解决先切回纯 YOLOv8 的默认结构把额外模块全部注释掉跑一个 baseline。如果 baseline 的 mAP 比花哨结构更高说明你的数据量撑不起这些模块。反过来如果 baseline mAP 也低就先解决数据质量问题再考虑结构优化。我现在的习惯是结构改动永远在 baseline 之后。5.3 显存看似充足却 OOM缓存策略和 batch 的配合现象nvidia-smi 看显存还有剩余但训练到某个 epoch 突然报 CUDA out of memory。原因YOLOv8 训练时默认会缓存图片到显存来加速数据读取。如果开了 cache 且输入分辨率是 640一张图的缓存占用的显存远大于单次前向的显存占用。我还有一次是开了随机裁剪增强某个 batch 里恰好出现几张尺寸最大的图峰值显存顶穿了。解决看训练日志里的 Batch 大小和时间先按 60% 显存占用量估算 batch。如果必须要用大 batch把 cache 关掉或者改成缓存到内存。代码里一般有个 cache 参数设成 False 就不会把整张图塞进显存。这个参数在 train.py 的 DataLoader 配置里看到 cacheTrue 就改成 False。5.4 预训练权重加载报 shape 错误类别数变化带来的维度不匹配现象运行 train.py 加载官方 yolov8s.pt 时报错说某个参数的 shape 对不上比如 Expected size 80 but got size 1。原因官方的预训练权重是在 COCO 数据集上训练的分类头输出 80 个类别。你的数据集只有 1 类分类头的卷积核维度是 80 不是 1直接加载必然报错。解决用 parse_weight.py 把预训练权重转换一遍。转换时只保留 backbone 和 neck 部分分类头和检测头用随机初始化。如果转换后训练效果还不如从头训练检查一下是不是 torch.load 时把整个 checkpoint 结构都加载进去了而不是只提取 state_dict。我当时就是这么翻车的后来强制在加载后打印一层参数的 shape 做验证才发现问题。5.5 训练中断只能从头再来没有预留恢复机制现象训练到 80 个 epoch机器重启或者显存被其他任务占了再启动训练只能从 epoch 0 开始。原因没走 resume 逻辑。train.py 一般在每个 epoch 结束保存一个 last.pt但它不会自动从 last.pt 续训。你需要手动在启动命令里指定 resume 路径。解决启动命令加一行--resume ./runs/exp/weights/last.pt。为了避免这种悲剧我后来每次训练都在脚本里先检查是否存在 last.pt存在就直接走 resume不存在才从头训练。这个习惯至少救过我三次每次都省下十几个小时的重复训练时间。6. 进阶把一次训练变成可回溯的完整实验6.1 损失曲线与置信度分布一起看训练结束别急着看 mAP先画损失曲线。train.py 会把每个 epoch 的 loss 值打印到日志或保存成数组文件读取后画图import matplotlib.pyplot as plt epochs list(range(1, len(loss_history) 1)) fig, ax plt.subplots(2, 1, figsize(10, 8)) ax[0].plot(epochs, loss_history[train_loss], labeltrain_loss) ax[0].plot(epochs, loss_history[val_loss], labelval_loss) ax[0].set_xlabel(epoch) ax[0].set_ylabel(loss) ax[0].legend() ax[1].plot(epochs, loss_history[mAP50], labelmAP50, colororange) ax[1].set_xlabel(epoch) ax[1].set_ylabel(mAP50) ax[1].legend() plt.tight_layout() plt.savefig(training_curve.png, dpi200)画完以后看两个点train_loss 和 val_loss 的开口。如果 train_loss 持续下降但 val_loss 在某个 epoch 后反弹就是过拟合信号把 epochs 减少或增强数据增强。mAP50 曲线如果在前半程快速上升、后半程走平说明模型容量和数据量基本匹配继续加 epochs 意义不大。损失曲线是被低估的工具它能告诉你模型到底是在正常收敛还是早就在硬蹭。6.2 用 predict.py 做误检筛查训练完以后把验证集图片跑一遍 predict.py关掉 nms 的显示阈值保留所有置信度大于 0.05 的框导出成带坐标和置信度的 txt。然后你按置信度从低到高翻图会看到两种典型情况置信度低但框得很准的说明训练不充分置信度很高但框到背景上的说明数据里背景样本太单一。我以前只用 mAP 评价模型直到有一次发现 mAP 0.72 但某个特定场景全部误检才意识到指标不能替代样本级检查。从那以后我每次训完都会强制走一遍低置信度框筛查把典型误检图回标到训练集里做 hard example 补充——这是我实践下来对精度提升最直接的一招希望帮到你。本文还有配套的精品资源点击获取