YOLOv5-6.1 单通道图像训练保姆级教程：从修改代码到成功导出ONNX模型

YOLOv5-6.1 单通道图像训练全流程实战指南

在工业质检、医学影像分析等专业场景中，灰度图像处理往往比RGB三通道图像更具实际意义。本文将手把手带你完成YOLOv5-6.1框架的单通道适配改造，从代码修改到模型导出，覆盖90%开发者会遇到的典型报错场景。不同于零散的代码片段分享，我们特别整理了全链路修改清单与报错解决方案对照表，确保即使没有深度学习经验的工程师也能顺利完成项目部署。

1. 环境准备与项目初始化

1.1 基础环境配置

推荐使用Python 3.8+和PyTorch 1.8+的组合，这是经过验证与YOLOv5-6.1兼容性最好的版本环境。使用conda创建隔离环境：

conda create -n yolov5_mono python=3.8 conda activate yolov5_mono pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

克隆官方仓库时建议指定6.1版本分支：

git clone -b v6.1 https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt

1.2 数据集结构调整

单通道训练数据集需要特殊处理目录结构。建议采用如下组织形式：

dataset_mono/ ├── images/ │ ├── train/ │ │ ├── image1.png │ │ └── image2.png │ └── val/ │ ├── image3.png │ └── image4.png └── labels/ ├── train/ │ ├── image1.txt │ └── image2.txt └── val/ ├── image3.txt └── image4.txt

注意：所有图像文件需统一为单通道格式（如PNG灰度图），避免混合通道数导致训练异常

2. 核心代码修改全解析

2.1 训练入口文件改造

train.py需要修改三处关键配置：

1. 默认通道数设置（约第440行）：

parser.add_argument('--channels', type=int, default=1, help='number of input channels')

2. 模型初始化部分（约第480行）：

model = Model(opt.cfg, ch=1, nc=nc, anchors=hyp.get('anchors')).to(device) # 修改ch=3为ch=1

3. 数据加载参数传递：

train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, opt, hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, workers=opt.workers, image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '), shuffle=True, channels=1) # 新增channels参数

2.2 数据集加载逻辑重构

utils/datasets.py需要系统性改造：

1. 图像读取方式修改（约第150行）：

img = cv2.imread(path, cv2.IMREAD_GRAYSCALE) # 替换原cv2.imread(path)

2. 通道维度处理（约第180行）：

if len(img.shape) == 2: img = np.expand_dims(img, axis=2) # 灰度图增加通道维度

3. 移除HSV增强（约第220行）：

# 注释掉以下HSV增强代码块 # if random.random() < hgain: # augment_hsv(img, hgain=hgain, sgain=sgain, vgain=vgain)

2.3 模型架构适配调整

models/yolo.py和models/common.py需要协同修改：

1. 输入通道验证（common.py约第320行）：

if im.mode != 'L': im = im.convert('L') # 强制转换为灰度模式

2. 卷积层权重校验（yolo.py约第40行）：

def __init__(self, cfg='yolov5s.yaml', ch=1, nc=None, anchors=None): # 修改默认ch=3为ch=1

3. 特征图处理（yolo.py约第120行）：

if m.f == -1: # 当来自上一层时 m.f = i - 1 # 修改特征图索引 m.in_channels = 1 # 单通道输入

3. 典型报错与解决方案

3.1 通道维度不匹配问题

错误现象：

RuntimeError: Given groups=1, weight of size [32, 3, 6, 6], expected input[8, 1, 640, 640] to have 3 channels

解决步骤：

1. 检查train.py中的--cfg参数是否指向正确配置文件
2. 确认models/yolov5s.yaml中nc和ch参数
3. 验证数据集加载是否真正输出单通道图像

3.2 多线程加载异常

错误现象：

AttributeError: 'NoneType' object has no attribute 'python_exit_status'

解决方案：

1. 修改train.py中workers参数：

parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers')

2. 添加环境变量（在文件开头插入）：

import os os.environ['KMP_DUPLICATE_LIB_OK']='True'

3.3 显存不足处理

当遇到CUDA out of memory时，可通过以下策略缓解：

4. ONNX模型导出实战

4.1 导出参数配置

修改export.py关键参数：

# 第457行附近修改 torch.onnx.export(model, im, f, verbose=False, opset_version=12, input_names=['images'], output_names=['output'], dynamic_axes={'images': {0: 'batch'}, # batch size可变 'output': {0: 'batch'}} if opt.dynamic else None, training=torch.onnx.TrainingMode.EVAL, do_constant_folding=True, export_params=True, keep_initializers_as_inputs=False, channels=1) # 新增单通道声明

4.2 导出验证脚本

使用以下代码验证ONNX模型：

import onnxruntime as ort import numpy as np sess = ort.InferenceSession("yolov5s_mono.onnx") input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name # 生成单通道测试数据 test_img = np.random.rand(1, 1, 640, 640).astype(np.float32) pred = sess.run([output_name], {input_name: test_img}) print(pred[0].shape) # 应输出(1, 25200, 85)

4.3 部署优化建议

1. 量化压缩：使用ONNX Runtime的量化工具减小模型体积

python -m onnxruntime.tools.convert_onnx_models_to_ort yolov5s_mono.onnx

2. 多后端支持：测试不同推理引擎的兼容性

# TensorRT测试 import tensorrt as trt

3. 性能对比：记录各平台推理耗时