YoloX训练实战：从零开始用PyTorch训练一个自定义数据集（附完整代码）

YoloX实战指南：PyTorch自定义数据集训练全流程解析

在工业检测、自动驾驶、安防监控等领域，目标检测技术正发挥着越来越重要的作用。YoloX作为Yolo系列的最新演进版本，凭借其Anchor-Free设计、SimOTA动态匹配等创新特性，在精度和速度上实现了显著提升。本文将带您从零开始，完成一个完整的YoloX模型训练流程，涵盖数据准备、模型配置、训练优化到结果评估的全过程。

1. 环境准备与数据标注

1.1 基础环境配置

YoloX训练需要PyTorch环境的支持，建议使用Python 3.8+和PyTorch 1.8+版本。以下是推荐的环境配置步骤：

# 创建conda环境 conda create -n yolox python=3.8 conda activate yolox # 安装PyTorch pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html # 安装YoloX依赖 pip install yolox==0.3.0 pip install opencv-python pillow tqdm tensorboard

对于硬件配置，建议至少使用8GB显存的GPU（如NVIDIA RTX 2070及以上）。如果显存不足，可以通过减小batch_size或输入图像尺寸来降低显存消耗。

1.2 数据集标注规范

YoloX支持多种标注格式，最常见的是VOC格式和COCO格式。本文以VOC格式为例，介绍数据准备过程：

1. 图像文件应存放在VOCdevkit/VOC2007/JPEGImages目录下
2. 标注文件（XML格式）存放在VOCdevkit/VOC2007/Annotations目录下
3. 每个XML文件对应一张图像，包含以下关键信息：

<annotation> <object> <name>person</name> <bndbox> <xmin>100</xmin> <ymin>200</ymin> <xmax>300</xmax> <ymax>400</ymax> </bndbox> </object> </annotation>

对于标注工具，推荐使用LabelImg或CVAT：

• LabelImg：轻量级开源工具，适合小规模标注

  pip install labelImg labelImg

• CVAT：功能强大的在线标注系统，支持团队协作

标注注意事项：确保标注框紧密贴合目标边缘，避免包含过多背景；对于遮挡目标，按可见部分标注；同类目标在不同场景下应保持标注一致性。

2. 数据预处理与增强

2.1 数据集划分与格式转换

使用voc_annotation.py脚本将数据集划分为训练集和验证集，并转换为YoloX需要的格式：

# voc_annotation.py关键参数配置 classes_path = 'model_data/custom_classes.txt' # 自定义类别文件 trainval_percent = 0.9 # 训练集+验证集比例 train_percent = 0.9 # 训练集比例

生成的train.txt文件每行格式为：

图像路径 x1,y1,x2,y2,class_id x1,y1,x2,y2,class_id ...

2.2 数据增强策略

YoloX默认使用了Mosaic和MixUp等增强技术，可以有效提升模型泛化能力。在yolox/data/data_augment.py中可以看到完整的增强流程：

1. Mosaic增强：随机选取4张图像拼接为1张
2. MixUp增强：两幅图像线性叠加
3. 色彩扰动：HSV空间随机调整色调、饱和度和明度
4. 几何变换：随机翻转、旋转、裁剪

对于特定场景（如工业缺陷检测），可能需要调整增强策略：

# 自定义增强配置示例 train_dataset = YoloXDataset( ... augment=True, degrees=10.0, # 旋转角度范围 translate=0.1, # 平移比例 scale=(0.5, 1.5), # 缩放范围 shear=2.0, # 剪切幅度 perspective=0.0, # 透视变换 )

2.3 处理类别不平衡问题

当数据集中各类别样本数量差异较大时，可采用以下策略：

1. 重采样：对少数类样本进行过采样
2. 损失加权：在YOLOLoss中调整各类别的权重

   class YOLOLoss(nn.Module): def __init__(self, num_classes, class_weights=None): if class_weights is not None: self.class_weights = torch.FloatTensor(class_weights)

3. 数据增强侧重：对少数类样本应用更强的增强

3. 模型配置与训练

3.1 模型选择与参数配置

YoloX提供多种规模的预训练模型，可根据硬件条件选择：

在train.py中关键配置参数：

# 模型选择 phi = 's' # 可选 'nano', 'tiny', 's', 'm', 'l', 'x' # 训练参数 input_shape = [640, 640] # 输入图像尺寸 batch_size = 8 # 根据显存调整 lr = 1e-3 # 初始学习率 num_workers = 4 # 数据加载线程数 # 训练策略 freeze_backbone = True # 是否冻结骨干网络 mosaic_prob = 1.0 # Mosaic增强概率 mixup_prob = 1.0 # MixUp增强概率

3.2 训练过程优化

YoloX采用两阶段训练策略：

1. 冻结阶段：冻结骨干网络（CSPDarknet），仅训练检测头

   # 冻结骨干网络 for param in model.backbone.parameters(): param.requires_grad = False

2. 解冻阶段：解冻全部网络，进行端到端微调

训练过程中可通过TensorBoard监控关键指标：

tensorboard --logdir=logs

常见训练问题与解决方案：

3.3 自定义模型结构修改

如需修改网络结构，可编辑yolox/models/yolo_pafpn.py和yolox/models/yolo_head.py。例如，增加一个检测层：

class YOLOXHead(nn.Module): def __init__(self, num_classes=80, strides=[8, 16, 32, 64]): # 新增64 stride super().__init__() # 新增一个检测头的配置 self.stems.append(BaseConv(in_channels, out_channels, ksize=1)) self.cls_convs.append(...) self.reg_convs.append(...)

4. 模型评估与部署

4.1 评估指标解读

YoloX训练过程中会输出多种评估指标：

• mAP(mean Average Precision)：综合衡量检测精度
• AP50：IoU阈值为0.5时的AP值
• AP75：IoU阈值为0.75时的AP值
• AR(Average Recall)：平均召回率

使用eval.py进行详细评估：

python tools/eval.py -n yolox-s -c yolox_s.pth -b 8 -d 1 --conf 0.001

4.2 模型导出与部署

YoloX支持导出为多种格式：

1. PyTorch转ONNX：

   python tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth

2. TensorRT加速：

   trtexec --onnx=yolox_s.onnx --saveEngine=yolox_s.engine --fp16

3. 移动端部署：
   • Android: 使用NCNN或MNN框架
   • iOS: 使用CoreML转换工具

4.3 实际应用技巧

1. 推理优化：

   # 设置更高效的推理参数 predictor = Predictor( model, exp, trt_file=None, decoder=None, device="gpu", fp16=False, legacy=False, )

2. 后处理调优：
   • 调整置信度阈值--conf（默认0.01）
   • 调整NMS阈值--nms（默认0.65）
3. 多尺度测试：

   exp.test_size = (736, 1280) # 使用更大尺寸测试提升小目标检测

5. 进阶技巧与问题排查

5.1 工业场景优化策略

针对工业检测的特殊需求，可采取以下优化措施：

1. 小目标检测优化：

   • 增加高分辨率检测层
   • 使用更密集的特征金字塔
   • 添加注意力机制

2. 类别不平衡处理：

   # 在SimOTA中调整类别权重 cost = cls_cost + 3.0 * iou_cost + 100000.0 * (~in_boxes_and_center).float()

3. 模型轻量化：

   • 使用通道剪枝
   • 知识蒸馏
   • 量化感知训练

5.2 常见错误排查

1. CUDA内存不足：

   • 减小batch_size
   • 使用--fp16混合精度训练
   • 尝试梯度累积：

     optimizer.zero_grad() for i, (images, targets) in enumerate(train_loader): loss = model(images, targets) loss.backward() if (i+1) % 4 == 0: # 每4步更新一次 optimizer.step() optimizer.zero_grad()

2. 训练出现NaN：

   • 检查数据中是否存在无效标注
   • 降低学习率
   • 添加梯度裁剪：

     torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)

3. 验证集性能差：

   • 检查训练/验证数据分布是否一致
   • 尝试关闭Mosaic增强
   • 调整数据增强强度

5.3 性能调优记录

以下是一个典型的调优过程记录：

在实际项目中，建议从简单配置开始，逐步增加复杂度，每次修改只调整一个变量，准确记录每次修改的影响。