YOLOX实战:从零构建专属目标检测模型(保姆级教程)
1. 环境准备与YOLOX安装
第一次接触YOLOX时,我也被各种依赖和环境配置搞得头大。后来发现只要按步骤来,其实半小时就能搞定。先确保你的机器有NVIDIA显卡(建议显存≥8GB),我用的是一台RTX 3060笔记本,实测完全够用。
1.1 基础环境搭建
推荐使用conda创建虚拟环境,避免污染系统环境。这里我用的是Python 3.8:
conda create -n yolox python=3.8
conda activate yolox
接着安装PyTorch 1.7+版本。注意要选对应CUDA版本的(我用的是CUDA 11.1):
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
1.2 获取YOLOX源码
官方代码库更新很频繁,建议用最新版:
git clone https://github.com/Megvii-BaseDetection/YOLOX
cd YOLOX
安装依赖时有个坑要注意:requirements.txt里的opencv-python版本可能冲突。我手动指定了4.5.4版本:
pip install -r requirements.txt opencv-python==4.5.4.60
1.3 编译扩展模块
YOLOX需要编译C++扩展,这个步骤容易出错。先确保g++版本≥5.0:
python setup.py develop
如果遇到apex安装失败(特别是Windows用户),可以跳过这个非必需组件。我在Ubuntu下用这个命令成功了:
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
2. 数据集准备与处理
去年做工业质检项目时,我收集了2000张带缺陷的电路板图片。这里分享下我的数据集处理经验。
2.1 数据格式选择
YOLOX支持VOC和COCO格式。对于新手我推荐VOC格式,结构更直观:
dataset/
├── Annotations/ # XML标注文件
├── JPEGImages/ # 原始图片
└── ImageSets/
└── Main/ # 包含train.txt, val.txt等
2.2 标注文件转换
如果你的数据是YOLO格式(txt标注),可以用这个Python脚本转换:
import xml.etree.ElementTree as ET
from glob import glob
import os
def yolo_to_voc(img_dir, txt_dir, output_dir):
os.makedirs(output_dir, exist_ok=True)
for txt_file in glob(f"{txt_dir}/*.txt"):
# 转换逻辑...
tree.write(f"{output_dir}/{base}.xml")
2.3 数据集划分
用sklearn的train_test_split快速划分训练集和验证集:
from sklearn.model_selection import train_test_split
all_images = sorted(glob("JPEGImages/*.jpg"))
train, val = train_test_split(all_images, test_size=0.2, random_state=42)
with open("ImageSets/Main/train.txt", "w") as f:
f.write("\n".join([os.path.splitext(os.path.basename(p))[0] for p in train]))
3. 关键代码修改实战
第一次修改代码时我踩了不少坑,这里把关键点都标注出来。
3.1 类别定义修改
修改yolox/data/datasets/voc_classes.py时要注意:
- 类别名必须用英文
- 最后一个类别后面不能加逗号
- 顺序要和标注文件一致
VOC_CLASSES = (
"crack",
"short_circuit" # 注意这里没有逗号
)
3.2 配置文件调整
在exps/example/yolox_voc/yolox_voc_s.py中修改:
class Exp(MyExp):
def __init__(self):
super(Exp, self).__init__()
self.num_classes = 2 # 改成你的类别数
self.depth = 0.33 # 小模型用这个参数
self.width = 0.50
self.data_dir = "/path/to/your/dataset" # 绝对路径!
self.train_ann = "train.txt"
self.val_ann = "val.txt"
self.max_epoch = 300 # 工业场景建议300+
3.3 数据增强配置
在yolox/data/data_augment.py中可以调整Mosaic和MixUp的概率。对于小数据集(<1万张),我建议:
self.degrees = 10.0 # 旋转角度范围
self.translate = 0.1 # 平移比例
self.scale = (0.5, 1.5) # 缩放范围
self.mosaic_prob = 0.8 # 工业缺陷检测可以调低到0.5
self.mixup_prob = 0.2 # 小数据建议0.3
4. 模型训练与调优
4.1 启动训练命令
单卡训练用这个命令(记得先下载预训练权重):
python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py \
-d 0 -b 32 -c yolox_s.pth.tar \
--fp16 -o # 启用混合精度训练
参数说明:
-d 0使用第0号GPU-b 32batch size根据显存调整--fp16能节省30%显存-o使用优化器动量
4.2 学习率调整策略
在yolox/exp/yolox_base.py中修改:
self.warmup_epochs = 5 # 小数据集可以增加到10
self.basic_lr_per_img = 0.01 / 64.0 # 学习率与batch size线性相关
self.no_aug_epochs = 20 # 最后20个epoch关闭数据增强
4.3 训练监控技巧
推荐用TensorBoard监控训练过程:
tensorboard --logdir=YOLOX_outputs/yolox_voc_s --bind_all
重点关注这三个指标:
- train/loss - 应该平稳下降
- val/AP - 验证集精度
- lr - 学习率变化曲线
5. 模型测试与部署
5.1 测试脚本使用
用demo.py快速验证效果:
python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s.py \
-c YOLOX_outputs/yolox_voc_s/latest_ckpt.pth \
--path assets/dog.jpg \
--conf 0.3 --nms 0.5 --tsize 640 \
--save_result --device gpu
5.2 量化部署方案
用TensorRT加速推理(速度提升3倍+):
from yolox.utils import trt_engine
engine = trt_engine.build_engine(
"yolox_s.onnx",
precision="fp16" # 也可以用int8量化
)
detector = trt_engine.YOLOXTRT(engine)
5.3 常见问题解决
-
CUDA out of memory:
- 减小batch size
- 添加
--fp16参数 - 修改
--tsize为更小的尺寸
-
验证集AP不升反降:
- 检查数据标注质量
- 降低学习率
- 增加
no_aug_epochs
-
预测框偏移严重:
- 检查数据增强参数
- 确认输入尺寸和训练时一致
- 验证标注框是否超出图像边界
训练完成后,建议用PyQt或Gradio快速搭建演示界面。我常用的Gradio demo代码不到50行就能实现实时检测效果展示。
更多推荐


所有评论(0)