YOLOv5模型训练翻车实录:从Ubuntu20.04环境配置到Pillow版本冲突的避坑指南
YOLOv5模型训练实战避坑指南:从环境搭建到版本冲突解决
第一次在Ubuntu系统上训练YOLOv5模型时,我遇到了无数令人抓狂的问题。从CUDA版本不兼容到Pillow库的诡异报错,每一步都像在拆解一个技术炸弹。本文将分享我在Ubuntu 20.04上使用ROS Noetic和CUDA 12.2环境训练交通信号灯检测模型时踩过的坑,以及如何系统性地解决这些问题。
1. 环境配置:那些看似简单却暗藏玄机的步骤
Ubuntu 20.04作为深度学习开发的黄金标准系统,理论上应该能提供稳定的基础环境。但当我开始搭建YOLOv5训练环境时,才发现事情没那么简单。
1.1 CUDA与cuDNN的版本迷宫
安装NVIDIA驱动和CUDA工具包时,第一个坑就出现了。官方文档推荐使用CUDA 11.3,但我的RTX 3080显卡需要CUDA 12.2才能发挥最佳性能。经过多次尝试,最终确定了这套组合:
# 安装NVIDIA驱动 sudo apt install nvidia-driver-535 # 安装CUDA 12.2 wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600 sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /" sudo apt-get update sudo apt-get -y install cuda-12-2注意:安装完成后务必验证CUDA版本是否与系统环境匹配,使用
nvcc --version检查
1.2 Conda环境管理的艺术
为了避免污染系统Python环境,我创建了专门的conda环境:
conda create -n yolov5 python=3.8 conda activate yolov5 pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116这里遇到了第二个坑:PyTorch与CUDA版本的兼容性问题。经过多次尝试,发现PyTorch 1.12.1与CUDA 11.6的组合最为稳定,尽管我安装了CUDA 12.2。
2. 数据准备:从标注到格式转换的完整流程
数据是模型训练的核心,但数据处理过程往往比想象中复杂得多。
2.1 LabelMe标注实战技巧
使用LabelMe进行标注时,有几个关键点需要注意:
- 标注顺序一致性:始终按顺时针或逆时针方向标注,确保多边形顶点顺序一致
- 标签命名规范:采用全小写、无空格的命名方式,如"red_light"而非"Red Light"
- 图像预处理:标注前统一调整图像尺寸,建议使用640x640分辨率
标注完成后,文件结构应该如下:
yolov5_train/ ├── images/ # 原始图像 ├── jsons/ # LabelMe生成的JSON标注文件 ├── txts/ # 转换后的YOLO格式标注 └── train_data/ # 最终训练数据集2.2 标注格式转换的关键脚本
将LabelMe的JSON格式转换为YOLO需要的TXT格式是个技术活。以下是核心转换代码:
def convert_labelme_to_yolo(json_file, output_dir, class_list): with open(json_file) as f: data = json.load(f) image_width = data['imageWidth'] image_height = data['imageHeight'] output_lines = [] for shape in data['shapes']: label = shape['label'] class_id = class_list.index(label) # 转换多边形坐标为YOLO格式 points = np.array(shape['points']) normalized_points = points / [image_width, image_height] flattened_points = normalized_points.flatten().tolist() line = [str(class_id)] + [str(round(p, 6)) for p in flattened_points] output_lines.append(' '.join(line) + '\n') output_path = os.path.join(output_dir, os.path.splitext(os.path.basename(json_file))[0] + '.txt') with open(output_path, 'w') as f: f.writelines(output_lines)提示:转换后务必检查生成的TXT文件,确认坐标值在0-1范围内且类别ID正确
3. 模型训练:那些官方文档没告诉你的细节
3.1 配置文件调优实战
YOLOv5的模型配置文件看似简单,实则暗藏玄机。以下是我的TLD.yaml配置关键部分:
# 数据配置 train: ../train_data/images/train/ val: ../train_data/images/valid/ test: ../train_data/images/test/ # 类别数量和名称 nc: 3 # red, green, yellow names: ['red', 'green', 'yellow'] # 超参数 lr0: 0.01 # 初始学习率 lrf: 0.2 # 最终学习率系数 momentum: 0.937 weight_decay: 0.0005 warmup_epochs: 3.0 warmup_momentum: 0.8 warmup_bias_lr: 0.13.2 Pillow版本冲突的终极解决方案
训练过程中最令人崩溃的错误莫过于:
ImportError: cannot import name 'PILLOW_VERSION' from 'PIL'这个问题源于新版本Pillow(>=10.0.0)移除了PILLOW_VERSION属性。解决方法有几种:
降级Pillow(推荐):
pip install Pillow==9.5.0修改YOLOv5源码: 在utils/datasets.py中找到相关导入,替换为:
try: from PIL import PILLOW_VERSION except ImportError: from PIL import __version__ as PILLOW_VERSION使用Docker环境:
docker pull ultralytics/yolov5:latest
经过测试,方案1最为简单可靠,不会引入其他兼容性问题。
4. 训练优化与结果分析
4.1 关键训练参数设置
启动训练时,这些参数对结果影响巨大:
python train.py \ --img 640 \ # 图像尺寸 --batch 16 \ # 批次大小 --epochs 100 \ # 训练轮次 --data data/TLD.yaml \ # 数据配置 --cfg models/TLD.yaml \ # 模型配置 --weights yolov5s.pt \ # 预训练权重 --cache \ # 启用缓存加速 --device 0 \ # 使用GPU 0 --adam \ # 使用Adam优化器 --single-cls # 单类别训练模式4.2 常见训练问题排查表
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| GPU利用率低 | 批次大小过小 | 增加--batch参数值 |
| 验证集mAP不升 | 学习率过高 | 降低--lr0参数值 |
| 训练损失震荡 | 数据不平衡 | 检查标注分布,增加数据增强 |
| 内存溢出 | 图像尺寸过大 | 减小--img参数值 |
| NaN损失 | 梯度爆炸 | 添加梯度裁剪,降低学习率 |
4.3 训练结果可视化技巧
训练完成后,使用YOLOv5内置工具分析结果:
python val.py \ --weights runs/train/exp/weights/best.pt \ --data data/TLD.yaml \ --task test \ --verbose关键指标解读:
- mAP@0.5: 交并比阈值为0.5时的平均精度
- mAP@0.5:0.95: 不同IOU阈值下的平均精度
- Precision: 预测为正样本中实际为正的比例
- Recall: 实际正样本中被正确预测的比例
5. 模型部署与性能优化
训练好的模型需要经过优化才能在实际应用中发挥价值。使用以下命令导出优化后的模型:
python export.py \ --weights runs/train/exp/weights/best.pt \ --img 640 \ --batch 1 \ --device 0 \ --include onnx \ --optimize导出ONNX格式后,可以使用TensorRT进一步加速:
trtexec \ --onnx=best.onnx \ --saveEngine=best.engine \ --fp16 \ --workspace=2048实际测试中,经过TensorRT优化的模型推理速度提升了3-5倍,显存占用减少40%。
