从零构建自定义数据集训练链路：mmclassification实战与避坑指南

1. 从零构建自定义数据集训练链路为什么值得认真对待搞图像分类的朋友大概率都经历过这样一个阶段拿开源框架跑通了官方Demo看着CIFAR-10或者ImageNet上的准确率蹭蹭往上涨心里挺美。但一旦要把自己手头那堆乱七八糟的图片喂进去训练立刻就卡壳了——文件夹怎么组织标注文件什么格式配置文件里那一堆参数到底改哪个数据增强怎么加预训练权重怎么接这一整套流程就是mmclassification自定义数据集并训练要解决的核心问题。它不是一个简单的“改个路径就行”的操作而是涉及数据组织规范、配置文件继承机制、训练策略调优、以及后续推理验证的完整链路。我见过太多人在这条链路上反复踩坑有人把图片和标签的对应关系搞反了导致训练loss不降有人改了配置文件但没生效因为不知道继承覆盖的优先级有人在自定义数据集上直接套用ImageNet的训练参数结果模型完全过拟合。这篇文章面向的是已经了解图像分类基本概念、跑过至少一个深度学习框架的Demo、但还没完整走过自定义数据训练流程的开发者。我会把整个链路拆开揉碎从数据准备到配置文件修改到训练启动到结果验证每一步都讲清楚“为什么这么做”而不只是“怎么做”。同时我会补充大量实际操作中才会遇到的细节问题比如类别不平衡怎么处理、小数据集怎么选增强策略、训练过程中怎么判断模型是否在正常收敛。这些内容在官方文档里往往一笔带过但实际做项目时恰恰是最耗时间的部分。整篇内容基于mmclassification这个框架来展开但其中的思路和方法论适用于绝大多数图像分类训练场景。无论你最终用的是哪个框架数据组织、配置管理、训练监控这些环节的底层逻辑是相通的。2. 数据准备自定义数据集的目录结构与标注规范2.1 为什么目录结构这么重要mmclassification对数据集的读取依赖于一套约定好的目录结构。很多人觉得“我把图片放一个文件夹里再写个txt记录标签不就行了”但框架在设计时考虑的是通用性和可扩展性它需要一种能够同时支持分类、检测、分割等多种任务的统一数据组织方式。具体到分类任务mmclassification支持两种主流的数据组织方式按类别分文件夹和基于标注文件的列表。前者适合数据量不大、类别清晰的场景后者适合数据来源复杂、需要灵活划分训练验证集的场景。我个人的建议是如果你的数据是从不同渠道收集来的或者需要精确控制每个类别的训练/验证比例用标注文件的方式更稳妥。2.2 按类别分文件夹的组织方式这是最直观的方式。假设你要做一个花卉分类任务有玫瑰、郁金香、向日葵三个类别目录结构应该是这样的flower_dataset/ ├── train/ │ ├── rose/ │ │ ├── img_001.jpg │ │ ├── img_002.jpg │ │ └── ... │ ├── tulip/ │ │ ├── img_001.jpg │ │ └── ... │ └── sunflower/ │ ├── img_001.jpg │ └── ... └── val/ ├── rose/ │ └── ... ├── tulip/ │ └── ... └── sunflower/ └── ...这里有几个容易出问题的地方。第一文件夹名称就是类别名称框架会按照文件夹名的字母顺序来分配类别索引。这意味着如果你的类别是“cat”和“dog”那cat的索引是0dog的索引是1。但如果你后面又加了一个“bird”类别索引顺序就变成了bird0, cat1, dog2之前训练好的模型标签就全乱了。所以一旦确定了类别集合就不要轻易增删类别。第二图片格式要统一。虽然框架底层用的PIL或者OpenCV能读取多种格式但混用jpg、png、bmp会导致读取速度不一致而且png的透明通道在某些预处理步骤中会引发意外错误。建议统一转成jpg格式色彩空间统一为RGB。第三图片尺寸不需要提前统一。很多人习惯先把所有图片resize到固定尺寸再喂给框架这其实多此一举。mmclassification的pipeline里包含了Resize和Crop操作你只需要在配置文件中指定目标尺寸即可。提前resize反而会损失信息尤其是当你后续想调整输入尺寸做对比实验时原始图片已经没了。2.3 基于标注文件的组织方式当数据量较大或者需要灵活划分数据集时标注文件的方式更合适。你需要准备一个txt文件每行包含图片路径和对应的类别索引train/rose/img_001.jpg 0 train/rose/img_002.jpg 0 train/tulip/img_001.jpg 1 train/sunflower/img_001.jpg 2 ...然后在配置文件的data部分指定这个txt文件的路径。这种方式的优势在于你可以精确控制哪些图片进入训练集、哪些进入验证集而不必受文件夹结构的限制。比如你可以把同一张图片的不同增强版本分别放入训练集和验证集或者按照拍摄时间划分训练验证集来模拟真实场景。注意标注文件中的路径是相对于配置文件中所设置的data_root的。很多人在这里搞混写成了绝对路径或者相对于项目根目录的路径导致框架找不到图片。建议统一使用相对路径并且在配置文件中明确设置data_root。2.4 数据集划分的实操建议训练集和验证集的划分比例没有绝对标准但有几个经验值可以参考。对于万级别以下的数据集通常按7:3或8:2划分对于十万级别以上的数据集9:1甚至95:5也够用。关键是要保证验证集的类别分布与训练集一致不能出现某个类别在验证集中完全没有样本的情况。如果数据量特别小比如每个类别只有几十张可以考虑使用交叉验证。mmclassification本身不直接支持交叉验证但你可以通过多次修改标注文件来手动实现。另一个技巧是使用分层抽样确保每个类别的训练验证比例一致。sklearn的train_test_split函数支持stratify参数可以很方便地实现这一点。还有一个容易被忽视的问题是数据泄漏。如果你是从视频中截取的帧那么相邻帧之间的差异非常小如果随机划分训练验证集很可能验证集中的图片和训练集中的图片几乎一样导致验证准确率虚高。这种情况下应该按视频来源划分同一个视频的帧要么全在训练集要么全在验证集。3. 配置文件体系继承机制与关键参数解析3.1 配置文件继承为什么不能直接改原始文件mmclassification的配置文件采用了一种继承机制这是很多人第一次接触时最困惑的地方。你打开一个模型配置文件比如resnet50_8xb32_in1k.py会发现里面只有寥寥几行_base_ [ ../_base_/models/resnet50.py, ../_base_/datasets/imagenet_bs32.py, ../_base_/schedules/imagenet_bs256.py, ../_base_/default_runtime.py ]这就是继承机制的核心基础配置放在_base_目录下模型配置文件通过_base_列表引入这些基础配置然后在此基础上做修改。这种设计的好处是避免重复比如所有基于ResNet的模型都共享同一套数据预处理流程只需要在各自的配置文件中修改模型结构相关的参数即可。但这也带来了一个问题修改的优先级。当你在模型配置文件中写了一个参数而基础配置中也有同名参数时以模型配置文件中的为准。但如果是嵌套字典比如要修改train_pipeline中的某个步骤就需要完整地重写整个train_pipeline列表而不是只写要改的那一项。这一点非常关键我见过不少人只写了要修改的那一行结果整个pipeline被覆盖成了只有那一行训练直接报错。3.2 自定义数据集配置的完整写法假设我们要基于ResNet50训练一个自定义的花卉分类数据集类别数为3。新建一个配置文件my_flower_resnet50.py内容如下_base_ [ ../_base_/models/resnet50.py, ../_base_/datasets/imagenet_bs32.py, ../_base_/schedules/imagenet_bs256.py, ../_base_/default_runtime.py ] # 修改类别数 model dict( headdict(num_classes3) ) # 修改数据集类型和路径 dataset_type CustomDataset data_root data/flower_dataset/ # 数据预处理pipeline train_pipeline [ dict(typeLoadImageFromFile), dict(typeRandomResizedCrop, scale224), dict(typeRandomFlip, prob0.5, directionhorizontal), dict(typePackInputs) ] test_pipeline [ dict(typeLoadImageFromFile), dict(typeResizeEdge, scale256, edgeshort), dict(typeCenterCrop, crop_size224), dict(typePackInputs) ] # 训练数据加载器 train_dataloader dict( batch_size32, num_workers4, datasetdict( typedataset_type, data_rootdata_root, ann_filetrain.txt, data_prefixtrain, pipelinetrain_pipeline ), samplerdict(typeDefaultSampler, shuffleTrue) ) # 验证数据加载器 val_dataloader dict( batch_size32, num_workers4, datasetdict( typedataset_type, data_rootdata_root, ann_fileval.txt, data_prefixval, pipelinetest_pipeline ), samplerdict(typeDefaultSampler, shuffleFalse) ) val_evaluator dict(typeAccuracy, topk(1,)) # 训练调度 optim_wrapper dict( optimizerdict(typeSGD, lr0.01, momentum0.9, weight_decay0.0001) ) param_scheduler dict( typeMultiStepLR, by_epochTrue, milestones[30, 60, 90], gamma0.1 ) train_cfg dict(by_epochTrue, max_epochs100, val_interval1) val_cfg dict() test_cfg dict() # 预训练权重 load_from https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth这个配置文件里每一块都有讲究。model部分的head.num_classes必须与你的实际类别数一致否则最后的全连接层维度对不上加载预训练权重时会报错。但这里有个技巧如果你设置了load_from框架会自动忽略全连接层的权重不匹配问题只加载骨干网络的权重。所以即使num_classes从1000改成了3预训练权重依然可以正常加载。3.3 学习率与批次大小的关系学习率不能随便设。有一个经验公式学习率与批次大小成正比。ImageNet上的标准配置是8卡×32批次256总批次学习率0.1。如果你只有一张卡批次大小设为32那么总批次是32学习率应该按比例缩小到0.1×32/2560.0125。我上面配置里写的0.01就是基于这个逻辑。但如果你用了预训练权重学习率可以适当大一点因为模型已经在一个大规模数据集上收敛到了一个不错的区域微调时可以用稍大的步长。反过来如果你的自定义数据集和预训练数据集差异很大比如ImageNet是自然图像你的数据是医学影像学习率反而要小一点避免把预训练学到的通用特征破坏掉。还有一个细节是权重衰减。SGD优化器通常配合weight_decay使用典型值是1e-4到5e-4。如果你的数据集很小过大的权重衰减会导致模型欠拟合如果数据集很大过小的权重衰减又容易过拟合。我一般从1e-4开始试观察训练loss和验证准确率的曲线来调整。3.4 数据增强策略的选择数据增强是提升模型泛化能力的关键手段但不是越多越好。对于自定义数据集增强策略的选择要基于数据的特点。花卉分类任务中水平翻转通常是安全的因为花的方向不影响类别判断。但垂直翻转就要谨慎因为自然界中花不会倒着长垂直翻转可能产生不合理的样本。RandomResizedCrop是一个很强的增强手段它随机裁剪图片的不同区域并缩放到目标尺寸。对于花卉数据这个操作可以让模型关注到花瓣、花蕊等局部特征通常能带来明显的性能提升。但要注意scale参数的设置默认是(0.08, 1.0)意味着最极端情况下只裁剪原图8%的区域。如果你的图片中目标占比很小这个参数可能需要调大下限。ColorJitter可以随机调整亮度、对比度、饱和度对于光照条件变化大的数据集很有帮助。但如果你的数据集本身就是在标准光照下采集的过度的颜色抖动反而会引入噪声。我一般会先跑一组不加ColorJitter的基线然后再加上看效果用数据说话。实操心得数据增强的效果高度依赖于具体数据集。不要盲目照搬ImageNet的增强配置建议先用较弱的增强跑一个基线然后逐步增加增强强度观察验证集准确率的变化。如果增加增强后验证准确率反而下降说明增强过度了。4. 训练启动与过程监控从命令到日志的完整解读4.1 训练命令与参数覆盖mmclassification的训练入口是tools/train.py最基本的命令是python tools/train.py my_flower_resnet50.py但实际使用中通常需要覆盖一些参数。比如你想把训练轮数从100改成50不需要去改配置文件直接在命令行覆盖即可python tools/train.py my_flower_resnet50.py --cfg-options train_cfg.max_epochs50这个--cfg-options参数非常实用它支持点号分隔的嵌套键值对。比如要修改学习率python tools/train.py my_flower_resnet50.py --cfg-options optim_wrapper.optimizer.lr0.005如果要指定工作目录和随机种子python tools/train.py my_flower_resnet50.py \ --work-dir work_dirs/flower_exp1 \ --seed 42工作目录会保存配置文件副本、日志文件、模型权重等。建议每次实验都指定不同的工作目录方便后续对比。随机种子的设置可以保证实验可复现这在调试阶段特别重要。4.2 日志解读如何判断训练是否正常训练启动后终端会输出类似这样的日志Epoch [1][50/100] lr: 1.000e-02 eta: 0:01:23 loss: 1.0986 Epoch [1][100/100] lr: 1.000e-02 eta: 0:00:00 loss: 1.0854 Epoch(val) [1][10/10] accuracy_top-1: 45.2300第一行中的loss是训练损失accuracy_top-1是验证集上的top-1准确率。判断训练是否正常主要看几个信号Loss是否在下降。第一个epoch的loss通常在ln(num_classes)附近比如3分类就是ln(3)≈1.0986。如果loss一直在这个值附近震荡不降说明模型没有在学习可能是学习率太小、数据有问题、或者标签和图片的对应关系错了。验证准确率是否在提升。如果训练loss在降但验证准确率不升反降说明模型在过拟合。这时候需要考虑增加数据增强、增大权重衰减、或者减少模型复杂度。学习率的变化。日志中的lr会随着scheduler变化。如果到了milestone轮数但lr没有下降说明scheduler配置有问题。4.3 训练过程中的常见异常与处理Loss变成NaN。这是最让人头疼的问题。常见原因有三个学习率太大导致梯度爆炸、数据中有损坏的图片、或者某个批次的标签超出了num_classes范围。排查方法是先用很小的学习率比如1e-5跑几个迭代如果loss正常下降说明是学习率问题如果还是NaN就要检查数据了。验证准确率始终在随机水平。比如3分类任务准确率一直在33%左右。这通常意味着模型完全没有学到东西。检查顺序是先确认标签是否正确可以手动打印几个样本的图片和标签看看再确认数据预处理是否合理比如归一化的mean和std是否匹配你的数据最后确认模型结构是否正确num_classes是否设置对了。训练速度异常慢。如果每个epoch耗时远超预期可能是num_workers设置得太小导致数据加载成为瓶颈或者图片分辨率太大导致前向传播计算量大。可以先用小分辨率图片跑通流程再逐步增大分辨率。显存溢出。减少batch_size是最直接的方法。另外可以开启混合精度训练在配置文件中加入optim_wrapper dict( typeAmpOptimWrapper, optimizerdict(typeSGD, lr0.01, momentum0.9, weight_decay0.0001) )混合精度训练可以在几乎不影响准确率的情况下减少约一半的显存占用。4.4 训练可视化与实验对比光看终端日志不够直观mmclassification支持对接TensorBoard等可视化工具。在配置文件中加入visualizer dict( typeVisualizer, vis_backends[ dict(typeLocalVisBackend), dict(typeTensorboardVisBackend) ] )训练启动后在另一个终端运行tensorboard --logdir work_dirs/flower_exp1就可以在浏览器中查看loss曲线、准确率曲线、学习率曲线等。多个实验的曲线可以叠加对比非常方便调参。实操心得我习惯在训练开始前先跑一个极小的子集比如每个类别10张图做冒烟测试确认整个流程能跑通、loss能下降、验证能执行。这个过程通常只需要几分钟但能避免跑了几个小时才发现配置有低级错误的情况。5. 模型验证与推理训练完之后怎么用5.1 在验证集上评估模型训练完成后用tools/test.py在验证集或测试集上评估模型python tools/test.py my_flower_resnet50.py work_dirs/flower_exp1/best_accuracy.pth输出会显示top-1和top-5准确率。如果你的数据集类别数少于5top-5准确率没有意义可以在配置文件中把val_evaluator的topk改成(1,)。评估时要注意一个问题验证集的预处理必须与训练时一致。具体来说训练时用了RandomResizedCrop和RandomFlip但验证时不能用这些随机增强只能用确定性的Resize和CenterCrop。这就是为什么配置文件中要分别定义train_pipeline和test_pipeline。如果验证时误用了训练pipeline每次评估结果都会不一样无法复现。5.2 单张图片推理实际使用中更常见的是对单张图片进行推理。mmclassification提供了tools/inference.py或者可以通过Python API调用from mmcls.apis import init_model, inference_model from mmcls.utils import register_all_modules register_all_modules() config_file my_flower_resnet50.py checkpoint_file work_dirs/flower_exp1/best_accuracy.pth model init_model(config_file, checkpoint_file, devicecuda:0) result inference_model(model, test_image.jpg) print(result.pred_class) print(result.pred_score)pred_class是预测的类别名称pred_score是各类别的置信度。如果置信度普遍偏低比如最高才0.4说明模型对这张图片不太确定可能是图片与训练数据分布差异较大或者模型本身训练不充分。5.3 模型导出与部署如果要把模型部署到生产环境通常需要导出为ONNX格式python tools/deployment/pytorch2onnx.py \ my_flower_resnet50.py \ work_dirs/flower_exp1/best_accuracy.pth \ --output-file flower_model.onnx \ --shape 1 3 224 224导出的ONNX模型可以用ONNXRuntime、TensorRT等推理引擎加载推理速度通常比PyTorch原生快不少。但要注意导出后的模型输入尺寸是固定的如果实际推理时图片尺寸不匹配需要先做resize。5.4 模型性能分析与优化方向训练完成后除了看准确率还应该分析模型在哪些类别上表现差。可以生成混淆矩阵from mmcls.apis import init_model, inference_model from sklearn.metrics import confusion_matrix import numpy as np # 收集所有验证集的预测结果和真实标签 preds [] labels [] for data in val_dataloader: result inference_model(model, data) preds.extend(result.pred_class) labels.extend(data.gt_label) cm confusion_matrix(labels, preds) print(cm)混淆矩阵能直观地看出哪些类别容易被混淆。比如玫瑰和郁金香如果互相误判很多说明模型没有学到区分这两个类别的关键特征。这时候可以考虑增加这两个类别的训练样本或者引入注意力机制让模型关注更细粒度的特征。另一个优化方向是类别不平衡处理。如果某个类别的样本数远少于其他类别模型会倾向于预测样本多的类别。解决方法包括对少数类别进行过采样、在损失函数中给少数类别更高的权重、或者使用focal loss等对不平衡数据更鲁棒的损失函数。6. 实操避坑指南与常见问题速查6.1 数据准备阶段的坑图片路径中有中文或空格。OpenCV读取图片时对中文路径支持不好可能导致读取失败。建议所有图片路径和文件名都使用英文和数字不要有空格和特殊字符。图片损坏或格式不支持。从网上下载的数据集经常有损坏的图片。可以在数据准备阶段写个脚本批量检查from PIL import Image import os for root, dirs, files in os.walk(data/flower_dataset): for f in files: path os.path.join(root, f) try: img Image.open(path) img.verify() except: print(f损坏图片: {path})类别文件夹名称大小写不一致。比如“Rose”和“rose”会被当成两个不同的类别。统一使用小写字母命名文件夹。6.2 配置阶段的坑num_classes忘记改。这是最常见的错误。加载预训练权重时如果num_classes不匹配框架会报错或者自动忽略全连接层。但如果你没设置load_from模型就会用随机初始化的全连接层训练效果完全取决于你的数据量。data_root路径写错。相对路径是相对于运行训练命令时的工作目录而不是配置文件所在的目录。建议在配置文件中使用绝对路径或者确保在项目根目录下运行训练命令。pipeline中的操作顺序。LoadImageFromFile必须放在第一位PackInputs必须放在最后一位。中间的操作顺序会影响增强效果比如先Crop再Flip和先Flip再Crop的结果是不一样的。6.3 训练阶段的坑显存不够。除了减小batch_size还可以尝试使用更小的输入分辨率、使用梯度累积通过OptimWrapper的accumulative_counts参数、使用更轻量的模型。训练loss震荡严重。可能是学习率太大尝试降低学习率或者使用warmup策略。在配置文件中加入param_scheduler [ dict(typeLinearLR, start_factor0.001, by_epochFalse, begin0, end500), dict(typeMultiStepLR, by_epochTrue, milestones[30, 60, 90], gamma0.1) ]验证准确率波动大。如果验证集太小准确率波动是正常的。可以增大验证集比例或者在验证时使用滑动窗口平均等方法平滑结果。6.4 常见问题速查表问题现象可能原因排查方法loss不下降学习率太小、数据标签错误、模型结构问题先用小学习率跑几个迭代检查标签对应关系loss变NaN学习率太大、数据有损坏、标签越界降低学习率检查数据完整性确认标签范围验证准确率随机水平模型没学到东西、预处理有问题检查num_classes、归一化参数、数据增强显存溢出batch_size太大、模型太大、分辨率太高减小batch_size使用混合精度降低分辨率训练速度慢num_workers太小、IO瓶颈、模型太大增大num_workers使用SSD存储换轻量模型过拟合数据量小、模型复杂、增强不足增加数据增强增大weight_decay使用早停类别不平衡某些类别样本太少过采样、类别权重、focal loss6.5 小数据集训练的特殊技巧当你的数据集很小比如每个类别只有几十张图片时直接从头训练几乎不可能得到好结果。这时候迁移学习是必选项。加载在ImageNet上预训练的权重冻结骨干网络只训练最后的全连接层。等全连接层收敛后再解冻部分骨干网络进行微调。冻结骨干网络的配置方法是在optim_wrapper中指定paramwise_cfgoptim_wrapper dict( optimizerdict(typeSGD, lr0.01, momentum0.9, weight_decay0.0001), paramwise_cfgdict( custom_keys{ backbone: dict(lr_mult0.1), # 骨干网络学习率降为十分之一 } ) )另一个技巧是使用更强的数据增强比如Mixup和CutMix。这两种增强方式通过混合两张图片及其标签来生成新的训练样本能显著提升小数据集上的泛化能力。在mmclassification中可以通过在pipeline中加入相应的transform来实现。实操心得小数据集训练时验证集的划分要特别小心。如果验证集太小准确率的方差会很大无法准确反映模型性能。我通常会用交叉验证来获得更稳定的评估结果虽然训练时间成倍增加但对于最终模型的选择更可靠。7. 从训练到落地一些个人体会整套流程走下来最深的体会是数据质量比模型结构重要得多。我见过太多人花大量时间调模型、换骨干网络、试各种trick但准确率提升有限。反而是花时间清洗数据、修正标注错误、补充难例样本之后准确率有了明显提升。如果你在自定义数据集上训练效果不理想先别急着换模型回头看看数据有没有问题。另一个体会是配置文件的管理。随着实验次数增多配置文件会越来越多很容易搞混哪个配置对应哪个结果。我的做法是每次实验都新建一个配置文件文件名包含日期和关键参数比如20240115_resnet50_lr001_bs32.py。同时在配置文件中用注释记录这次实验的目的和预期。这样过一段时间回头看能快速回忆起每次实验的意图。关于训练轮数很多人喜欢设一个很大的值然后靠早停来终止。但早停需要频繁验证而验证本身是有开销的。对于小数据集我通常先设50个epoch跑一次观察验证准确率在第几个epoch达到峰值然后把这个峰值附近的轮数作为后续实验的参考。这样比盲目设100个epoch要高效得多。最后说一个容易被忽视的点模型推理时的预处理必须与训练时完全一致。我遇到过好几次训练时准确率很高但部署后效果很差的情况排查后发现是推理时的归一化参数和训练时不一致。训练时用了ImageNet的mean和std推理时却用了默认值导致输入分布偏移。这种问题很隐蔽因为推理不会报错只是结果不准。建议把预处理参数写在一个公共的配置文件中训练和推理都从这里读取避免不一致。