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 32 batch 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

重点关注这三个指标:

  1. train/loss - 应该平稳下降
  2. val/AP - 验证集精度
  3. 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 常见问题解决

  1. CUDA out of memory

    • 减小batch size
    • 添加--fp16参数
    • 修改--tsize为更小的尺寸
  2. 验证集AP不升反降

    • 检查数据标注质量
    • 降低学习率
    • 增加no_aug_epochs
  3. 预测框偏移严重

    • 检查数据增强参数
    • 确认输入尺寸和训练时一致
    • 验证标注框是否超出图像边界

训练完成后,建议用PyQt或Gradio快速搭建演示界面。我常用的Gradio demo代码不到50行就能实现实时检测效果展示。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐