mmclassification自定义数据集训练全流程：从配置到部署

1. 从零上手 mmclassification自定义数据集训练全流程拆解做图像分类项目的人绕不开一个现实问题公开数据集跑通了换成自己的数据就各种报错。我最初接触 mmclassification 的时候也是被它的配置文件体系和数据集注册机制折腾了好一阵。但用熟了之后发现这套框架在自定义数据集训练这件事上其实设计得相当规整只是入门门槛稍微高了一点。mmclassification 是 OpenMMLab 生态里的图像分类工具箱基于 PyTorch 构建核心优势在于模块化设计——骨干网络、颈部、分类头、数据增强策略都可以通过配置文件灵活组合。它内置了大量经典模型ResNet、VGG、MobileNet、Vision Transformer 等和训练技巧Mixup、CutMix、Label Smoothing、Cosine Annealing 等适合想快速验证想法又不想从零搭训练框架的人。这篇文章面向两类读者一是刚接触 mmclassification、想用自己的图片数据跑通训练流程的开发者二是用过但被配置文件搞晕、想系统理清自定义数据集接入逻辑的从业者。我会从目录结构设计、标注文件生成、配置文件改写、训练启动与调试、常见报错排查几个维度把整条链路拆开讲清楚附带可直接复用的代码和配置片段。2. 自定义数据集接入的整体设计思路2.1 为什么 mmclassification 的数据集接入让人困惑mmclassification 的数据集体系继承自 MMDetection 的设计哲学采用注册器机制加配置文件驱动。好处是扩展性强坏处是新手不知道数据从哪进来、标签怎么对应、配置文件里哪些字段必须改。核心逻辑其实就三层数据集类负责读取标注文件并返回图像路径和标签数据流水线负责图像加载、增强、格式化数据加载器负责批处理和并行加速。自定义数据集的工作本质上就是让这三层都认识你的数据。2.2 两种主流接入方案的选择接入自定义数据集有两条路可走。第一条是复用内置的 CustomDataset 类只需要按照它约定的目录结构和标注格式组织数据然后写一个继承自该类的子类指定类别列表即可。第二条是完全自定义 Dataset 类适合标注格式特殊、需要复杂预处理逻辑的场景。对于绝大多数图像分类任务我强烈建议走第一条路。原因很简单CustomDataset 已经处理好了标注解析、类别映射、图像读取这些通用逻辑你只需要把数据摆放整齐代码量极少。只有当你需要从数据库读取图像、或者标注信息嵌在文件名里这类特殊情况才考虑完全自定义。2.3 目录结构设计的经验法则我见过太多人把数据随便丢在一个文件夹里就开始改配置结果训练时找不到图或者标签对不上。建议从一开始就按标准结构组织mydata/ ├── train/ │ ├── cat/ │ │ ├── 001.jpg │ │ └── 002.jpg │ ├── dog/ │ │ ├── 001.jpg │ │ └── 002.jpg │ └── bird/ │ └── ... ├── val/ │ ├── cat/ │ ├── dog/ │ └── bird/ └── test/ └── ...每个类别一个子文件夹文件夹名就是类别名。这种结构的好处是标注文件可以自动生成不需要手工写。另外训练集、验证集、测试集分开存放避免数据泄露。注意类别文件夹的命名不要用中文或特殊字符虽然 Python 能处理但在某些系统的文件编码下容易出问题。用英文小写加下划线最稳妥。3. 标注文件生成与数据集类编写实操3.1 自动生成标注文件的脚本mmclassification 支持的标注格式有两种txt 文件和 json 文件。txt 格式每行是“图像路径 标签”的形式json 格式则包含更丰富的元信息。我习惯用 txt简单直接。下面这个脚本可以扫描目录结构自动生成三个标注文件import os def generate_annotation(data_root, split, class_names): lines [] split_dir os.path.join(data_root, split) for label, cls_name in enumerate(class_names): cls_dir os.path.join(split_dir, cls_name) if not os.path.isdir(cls_dir): continue for fname in sorted(os.listdir(cls_dir)): if fname.lower().endswith((.jpg, .jpeg, .png, .bmp)): rel_path os.path.join(split, cls_name, fname) lines.append(f{rel_path} {label}) out_file os.path.join(data_root, f{split}.txt) with open(out_file, w, encodingutf-8) as f: f.write(\n.join(lines)) print(f{split}: {len(lines)} samples written to {out_file}) if __name__ __main__: data_root ./mydata classes [cat, dog, bird] for split in [train, val, test]: generate_annotation(data_root, split, classes)这里有个细节值得说明路径写的是相对路径相对于 data_root。这样整个数据集文件夹可以随意移动只要配置文件里的 data_root 跟着改就行不用重新生成标注文件。3.2 自定义 Dataset 类的编写有了标注文件接下来写数据集类。继承 CustomDataset只需要指定 classes 和 data_rootfrom mmcls.datasets import CustomDataset from mmcls.registry import DATASETS DATASETS.register_module() class MyDataset(CustomDataset): CLASSES (cat, dog, bird)就这么几行。CLASSES 元组的顺序必须和标注文件里的标签编号一致否则训练出来的模型会把猫认成狗。这个类可以放在项目目录下的 datasets 文件夹里然后在配置文件中通过 custom_imports 引入。如果你不想新建文件也可以直接在配置文件中用 typeCustomDataset 并传入 classes 参数效果一样dataset_type CustomDataset classes [cat, dog, bird]3.3 数据增强流水线的配置要点数据流水线决定了图像在送入网络前经历哪些变换。训练和验证阶段的流水线不同训练阶段需要随机增强来提升泛化能力验证阶段只需要确定性的预处理。一个典型的训练流水线配置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), ]RandomResizedCrop 的 scale224 表示随机裁剪并缩放到 224x224这是 ImageNet 系列模型的标准输入尺寸。如果你的模型输入尺寸不同这里要跟着改。ResizeEdge 加 CenterCrop 是验证阶段的标准做法保证每次评估结果一致。实操心得如果数据集较小每类少于 500 张建议加上 RandAugment 或 AutoAugment增强效果比手动堆叠变换更好。但要注意增强过强可能导致欠拟合需要根据验证集准确率调整。4. 配置文件改写与训练启动全流程4.1 配置文件继承机制的理解mmclassification 的配置文件采用继承机制基础配置定义了模型结构、训练策略、数据流水线等默认值你只需要在子配置中覆盖需要改的字段。这种设计避免了重复配置但也容易让人搞不清最终生效的值是什么。建议的做法是先选一个和你的任务最接近的官方配置作为基础比如 resnet18_8xb32_in1k.py然后新建一个配置文件继承它只改数据集相关和类别数相关的字段。_base_ [ ../_base_/models/resnet18.py, ../_base_/datasets/imagenet_bs32.py, ../_base_/schedules/imagenet_bs256.py, ../_base_/default_runtime.py, ] data_root ./mydata classes [cat, dog, bird] num_classes len(classes) model dict( headdict(num_classesnum_classes), ) train_dataloader dict( datasetdict( typeCustomDataset, data_rootdata_root, ann_filetrain.txt, data_prefix, classesclasses, pipelinetrain_pipeline, ), )这里有几个关键点需要解释。head 里的 num_classes 必须改成你的类别数否则最后的全连接层输出维度不对训练直接报错。data_prefix 设为空字符串是因为标注文件里的路径已经包含了 split 目录名如果标注文件里只有文件名这里就要设为对应的子目录。4.2 学习率与批次大小的适配计算官方配置的默认学习率是针对 8 卡、每卡 32 批次、共 256 总批次设定的。如果你只有单卡、批次大小设为 16学习率需要按线性缩放规则调整新学习率 基础学习率 × (新总批次 / 基础总批次)假设基础学习率是 0.1基础总批次 256你的总批次是 16那么新学习率 0.1 × 16 / 256 0.00625。这个规则在批次大小变化不大时比较准如果批次小到 8 以下建议用更保守的缩放或者配合梯度累积。optim_wrapper dict( optimizerdict(typeSGD, lr0.00625, momentum0.9, weight_decay0.0001), )4.3 训练启动命令与参数说明配置文件准备好后用 tools/train.py 启动训练python tools/train.py configs/mydataset/resnet18_my.py \ --work-dir work_dirs/resnet18_my \ --seed 42--work-dir 指定日志和权重保存目录--seed 固定随机种子方便复现。如果要从断点恢复训练加上 --resume 参数。想用预训练权重加速收敛在配置文件的 model 里加 init_cfgmodel dict( backbonedict( init_cfgdict( typePretrained, checkpointhttps://download.openmmlab.com/mmclassification/v0/resnet/resnet18_8xb32_in1k_20210831-fbbb1da6.pth, ), ), headdict(num_classesnum_classes), )预训练权重能显著加快小数据集的收敛速度通常能把训练轮数减少一半以上。4.4 训练过程监控与日志解读训练启动后终端会输出每个 epoch 的损失和准确率。重点关注几个指标训练损失是否稳定下降、验证准确率是否在提升、两者差距是否过大。如果训练损失降但验证准确率不升说明过拟合了需要加数据增强或正则化。如果训练损失震荡严重可能是学习率太大。日志文件保存在 work-dir 下的时间戳目录里可以用 tools/analysis_tools/analyze_logs.py 画曲线图直观判断训练状态。5. 常见报错与排查技巧实录5.1 数据集路径类报错最常见的报错是 FileNotFoundError提示找不到图像文件。排查思路先确认 data_root 和 ann_file 拼接后的路径是否正确再检查标注文件里的路径是相对路径还是绝对路径。如果标注文件里写的是绝对路径data_prefix 就不起作用了。另一个容易忽略的点是路径分隔符。Windows 下生成的标注文件用的是反斜杠在 Linux 环境读取时会出问题。生成标注文件时统一用 os.path.join 或者正斜杠能避免跨平台问题。5.2 类别数不匹配报错RuntimeError: size mismatch 通常是因为模型输出维度和数据集类别数不一致。检查两处配置文件中 model.head.num_classes 是否等于实际类别数数据集类的 CLASSES 长度是否和标注文件中的最大标签加一相等。5.3 显存不足的处理策略显存不够时优先降低批次大小其次减小输入尺寸。如果降到批次 4 还是不够可以启用梯度累积train_cfg dict( typeEpochBasedTrainLoop, max_epochs100, val_begin1, val_interval1, ) optim_wrapper dict( accumulative_counts4, )accumulative_counts4 表示每 4 个批次更新一次参数等效于批次大小扩大 4 倍但显存占用不变。5.4 常见问题速查表报错信息可能原因解决方法FileNotFoundError路径拼接错误检查 data_root 和 ann_filesize mismatch类别数不一致核对 num_classes 和 CLASSESCUDA out of memory批次过大降低批次或启用梯度累积KeyError: xxx配置字段名错误对照官方配置检查字段准确率不上升学习率不当按线性缩放调整学习率损失为 nan学习率过大降低学习率或加梯度裁剪避坑技巧每次改完配置文件先用 tools/misc/print_config.py 打印最终生效的完整配置确认所有字段都符合预期再启动训练。这个习惯能省下大量排查时间。6. 训练调优与模型评估的实战经验6.1 小数据集的迁移学习策略自定义数据集通常规模不大从零训练很难收敛。我的做法是分阶段训练先冻结骨干网络只训练分类头 5 到 10 个 epoch让分类头适应新类别然后解冻全部参数用较小的学习率微调。这样比直接端到端训练稳定得多。冻结骨干网络的配置方式model dict( backbonedict( frozen_stages4, init_cfgdict(typePretrained, checkpoint...), ), )frozen_stages4 表示冻结前四个 stageResNet18 一共四个 stage相当于整个骨干都冻住了。6.2 评估指标的选择与解读mmclassification 默认用 Top-1 和 Top-5 准确率。对于类别数少于 10 的任务Top-5 意义不大重点看 Top-1。如果类别不平衡还要关注每类的召回率避免模型偏向多数类。评估命令python tools/test.py configs/mydataset/resnet18_my.py \ work_dirs/resnet18_my/best_top1_acc_epoch_50.pth \ --metrics accuracy6.3 模型导出与部署准备训练完成后可以把模型导出为 ONNX 格式方便部署python tools/deployment/pytorch2onnx.py \ configs/mydataset/resnet18_my.py \ work_dirs/resnet18_my/best_top1_acc_epoch_50.pth \ --output-file resnet18_my.onnx \ --shape 1 3 224 224导出时注意输入 shape 要和训练时一致否则推理结果会偏差很大。我个人在实际操作中的体会是mmclassification 的自定义数据集接入最难的不是写代码而是理解它的配置继承逻辑和数据流走向。一旦把数据从标注文件到数据加载器这条链路理清楚后面换模型、换增强策略都是改几行配置的事。另外强烈建议第一次跑通时用少量数据快速验证流程确认无误再上全量数据能避免很多无效等待。