资讯动态

MMClassification自定义数据集实战:从Dataset到Pipeline的完整指南

发布时间:2026/10/18 1:20:27 来源:尧图企业网站定制
1. 从零开始理解mmclassification的数据集抽象层很多刚接触mmclassification的朋友第一反应是直接翻配置文件看到data_root、ann_file、classes这些字段就开始照猫画虎地改结果训练跑起来要么报KeyError要么精度死活上不去。我一开始也是这么干的后来踩了几次坑才明白mmclassification的数据集设计有一套自己的抽象逻辑不理解这套逻辑改配置就是盲人摸象。mmclassification把数据集拆成了两层Dataset类负责“怎么读一条样本”Data pipeline负责“读出来之后怎么变换”。这两层是解耦的同一个Dataset可以配不同的pipeline同一个pipeline也能套在不同的Dataset上。自定义数据集的核心工作其实就是让这两层都认识你的数据格式。1.1 为什么不能直接拿ImageFolder的思路来套如果你之前用PyTorch写过图像分类大概率用过torchvision.datasets.ImageFolder。它的逻辑很简单一个类别一个文件夹文件夹名就是类别名图片往里一扔就完事。这个思路在mmclassification里行不通原因有三个。第一mmclassification的Dataset返回的不是(image, label)这么简单它返回的是一个字典里面除了img和gt_label还可能带img_shape、ori_shape、flip、flip_direction这些元信息。这些元信息在训练时看着没用但到了测试阶段做多尺度推理、翻转增强的时候全靠它们来还原坐标和方向。第二mmclassification支持多标签分类一条样本可以对应多个标签。ImageFolder那种“一个文件夹一个类别”的硬结构根本表达不了这种关系。所以它的标注信息是存在一个独立的标注文件里的图片和标签之间通过文件名或索引来关联。第三mmclassification的pipeline是可插拔的数据增强、归一化、格式化都是独立的步骤。Dataset只负责把原始数据读进来至于怎么变换那是pipeline的事。这种设计的好处是你想换一种增强策略只需要改配置里的pipeline列表不用动Dataset代码。理解了这三点你就知道为什么自定义数据集不能简单粗暴地复制文件夹结构了。你得先想清楚你的标注信息以什么形式存在是txt、json还是csv图片路径和标签怎么对应想清楚这些再动手写代码。1.2 标注文件的三种常见组织方式在实际项目里我见过也用过三种标注组织方式各有各的适用场景。第一种是txt列表式每行一条样本格式是图片相对路径 标签。这种最直观用记事本就能改适合类别数不多、样本量中等的场景。缺点是路径里不能有空格标签只能是单个整数多标签就歇菜了。第二种是json式把所有样本组织成一个列表每个元素是一个字典包含img_path和gt_label字段。这种结构灵活多标签就是gt_label给个列表还能塞额外的元信息。缺点是文件大了之后读写慢几十万条样本的json动辄几百MB加载一次要等好几秒。第三种是文件夹加txt混合式图片按类别放在不同文件夹里但额外维护一个txt记录每个文件夹对应的类别索引。这种适合从ImageFolder迁移过来的项目改造成本低。缺点还是那个老问题多标签表达不了。我个人的建议是中小规模项目直接用txt列表式简单可靠调试的时候一眼就能看出问题。大规模或多标签项目用json式虽然加载慢点但结构清晰扩展性强。至于混合式除非你是从旧项目迁移否则没必要给自己找麻烦。2. 自定义Dataset类的编写要点与常见陷阱知道了数据怎么组织接下来就是写Dataset类。mmclassification里自定义Dataset有两种路子一种是继承BaseDataset另一种是继承CustomDataset。前者更底层需要自己实现load_data_list方法后者是官方封装好的支持从txt或文件夹读取改改配置就能用。我建议新手先从CustomDataset入手等熟悉了再往底层走。2.1 继承CustomDataset时最容易忽略的配置项CustomDataset的配置看起来简单但有几个字段特别容易踩坑。我列个表你对照着检查。配置项作用常见错误data_root数据集根目录写成绝对路径导致换机器就报错ann_file标注文件路径路径拼接时多一层或少一层目录data_prefix图片存放的子目录留空时默认在data_root下找图classes类别名称列表顺序和标注文件里的标签索引对不上test_mode是否测试模式测试时忘了关增强导致结果不稳定这里面最坑的是classes的顺序。mmclassification里标签是用整数索引表示的classes列表的第0个元素对应标签0第1个元素对应标签1以此类推。如果你标注文件里写的标签是1、2、3而classes列表只有三个元素那标签3就会越界报错。我见过有人把标签从1开始编号结果训练时loss一直不降查了半天才发现是标签偏移了一位。还有一个隐蔽的坑是data_prefix和ann_file里的路径拼接。CustomDataset在读取图片时会把data_root、data_prefix和标注文件里的相对路径拼在一起。如果你标注文件里写的是cat/001.jpgdata_prefix又设成了train那最终路径就是data_root/train/cat/001.jpg。多一层少一层都会导致FileNotFoundError。我的习惯是标注文件里只写文件名目录结构通过data_prefix来控制这样路径逻辑最清晰。2.2 重写load_data_list的完整示例如果你需要更灵活的控制比如从数据库读标注、或者图片路径需要动态计算那就得继承BaseDataset自己实现load_data_list。下面是一个完整的示例我加了详细注释。from mmcls.datasets import BaseDataset from mmcls.registry import DATASETS import os DATASETS.register_module() class MyDataset(BaseDataset): def __init__(self, data_root, ann_file, classes, **kwargs): self.data_root data_root self.ann_file ann_file self.classes classes super().__init__(**kwargs) def load_data_list(self): data_list [] with open(self.ann_file, r) as f: for line in f: line line.strip() if not line: continue parts line.split() img_path os.path.join(self.data_root, parts[0]) label int(parts[1]) data_list.append({ img_path: img_path, gt_label: label, }) return data_list这段代码看着简单但有几个细节值得说。第一DATASETS.register_module()这个装饰器必须加否则配置文件里用typeMyDataset的时候找不到类。第二load_data_list返回的每个字典里img_path必须是完整路径gt_label必须是整数或整数列表。第三super().__init__(**kwargs)要放在最后因为父类的初始化会调用load_data_list你得先把data_root这些属性准备好。注意如果你在load_data_list里做了耗时操作比如读大文件建议加个缓存机制否则每次__getitem__都重新读一遍训练速度会慢得让你怀疑人生。2.3 多标签场景下的标签处理多标签分类在mmclassification里是通过gt_label传列表来实现的。比如一条样本同时属于“猫”和“动物”两个类别gt_label就是[0, 5]。但这里有个坑多标签的损失函数和单标签不一样。单标签用CrossEntropyLoss多标签得用AsymmetricLoss或BinaryCrossEntropyLoss。如果你数据集是多标签的但配置里还用默认的交叉熵训练出来的模型会把多个标签当成互斥的精度直接崩掉。配置多标签损失的方法是在model的head里指定loss字段。比如headdict( typeMultiLabelLinearClsHead, num_classes10, in_channels512, lossdict(typeAsymmetricLoss, gamma_neg4, gamma_pos1), )另外多标签的标注文件里标签通常用空格分隔的多个整数表示比如img001.jpg 0 3 7。你在load_data_list里解析的时候要把它们转成列表而不是只取第一个。3. 数据增强pipeline的定制与调试数据集写好了接下来是pipeline。mmclassification的pipeline是一个列表里面每个元素是一个变换操作按顺序执行。训练时的典型pipeline长这样LoadImageFromFile→RandomResizedCrop→RandomFlip→PackInputs。看着简单但每个环节都有讲究。3.1 训练与测试pipeline的关键差异训练和测试的pipeline必须分开配置这是新手最容易犯的错误之一。训练时我们需要随机裁剪、随机翻转来增加数据多样性但测试时这些随机操作会导致同一个样本每次推理结果都不一样评估指标就没法复现了。训练pipeline的典型配置train_pipeline [ dict(typeLoadImageFromFile), dict(typeRandomResizedCrop, scale224), dict(typeRandomFlip, prob0.5, directionhorizontal), dict(typePackInputs), ]测试pipeline的典型配置test_pipeline [ dict(typeLoadImageFromFile), dict(typeResizeEdge, scale256, edgeshort), dict(typeCenterCrop, crop_size224), dict(typePackInputs), ]注意测试时用的是ResizeEdge加CenterCrop先把短边缩放到256再从中心裁224。这个组合是ImageNet时代的经典做法能保证输入尺寸固定同时保留图像主体内容。训练时的RandomResizedCrop则是随机裁一块再缩放裁出来的区域可能只包含物体的一部分这正是它增强效果的来源。还有一个细节是PackInputs必须放在pipeline的最后。它的作用是把前面变换得到的numpy数组打包成mmengine需要的格式包括转tensor、加batch维度等。如果你把它放在中间后面的变换操作就没法处理了。3.2 自定义增强操作的注册与使用mmclassification内置的增强操作已经覆盖了大部分场景但有时候你需要自己的变换逻辑。比如做医学图像分类时可能需要特定的窗宽窗位调整做遥感图像时可能需要旋转任意角度。这时候就得自定义变换类。自定义变换类需要继承BaseTransform并实现transform方法。下面是一个随机旋转任意角度的示例from mmcv.transforms import BaseTransform from mmcls.registry import TRANSFORMS import cv2 import numpy as np TRANSFORMS.register_module() class RandomRotate(BaseTransform): def __init__(self, max_angle30, prob0.5): self.max_angle max_angle self.prob prob def transform(self, results): if np.random.rand() self.prob: return results img results[img] h, w img.shape[:2] angle np.random.uniform(-self.max_angle, self.max_angle) center (w // 2, h // 2) matrix cv2.getRotationMatrix2D(center, angle, 1.0) results[img] cv2.warpAffine(img, matrix, (w, h)) return results注册好之后在配置文件的pipeline里加上dict(typeRandomRotate, max_angle45)就能用了。这里有个经验自定义变换尽量只操作results字典里的字段不要引入外部状态。因为mmclassification的dataloader是多进程的外部状态在进程间不共享容易出诡异的问题。3.3 pipeline调试的实用技巧pipeline出问题的时候报错信息往往很模糊比如KeyError: img_shape你根本不知道是哪个变换没设置这个字段。我的调试方法是写一个最小脚本手动跑一遍pipeline把中间结果打印出来。from mmcv.transforms import Compose pipeline Compose([ dict(typeLoadImageFromFile), dict(typeRandomResizedCrop, scale224), dict(typePackInputs), ]) results {img_path: test.jpg} results pipeline(results) print(results.keys()) print(results[img].shape)这样跑一遍你就能看到每个变换之后results里多了哪些字段、少了哪些字段。如果某个变换报错把pipeline截断到它前面单独跑那一个变换问题就定位了。提示调试pipeline时把num_workers设为0这样报错信息会直接打印在终端而不是被多进程吞掉。4. 配置文件组装与训练启动的完整流程数据集和pipeline都准备好了最后一步是把它们组装进配置文件然后启动训练。mmclassification的配置文件是继承式的你可以基于官方提供的resnet18_8xb32_in1k.py改也可以从零写一个。我倾向于从零写因为继承式配置改多了之后你根本不知道最终生效的是哪个值。4.1 配置文件的核心字段拆解一个完整的自定义数据集配置文件包含这几个部分data_preprocessor、model、train_dataloader、val_dataloader、optim_wrapper、param_scheduler、train_cfg、val_cfg、val_evaluator、default_hooks、work_dir。我挑几个和自定义数据集直接相关的说。data_preprocessor负责归一化和通道转换。ImageNet预训练模型通常用mean[123.675, 116.28, 103.53]、std[58.395, 57.12, 57.375]这是RGB顺序的均值和标准差。如果你的图片是BGR格式得先转成RGB否则归一化就错了。train_dataloader里的dataset字段指定数据集类型和参数sampler控制采样策略batch_size和num_workers影响训练速度。这里有个经验值num_workers设为CPU核心数的70%左右设太大了反而会因为进程切换开销导致速度下降。val_evaluator指定评估指标。单标签分类用Accuracy多标签用AveragePrecision。如果你自定义了数据集但评估指标没改跑出来的结果可能完全没意义。4.2 从单卡调试到多卡训练的平滑过渡新手建议先在单卡上把流程跑通确认loss能降、精度能涨再上多卡。单卡调试时把batch_size设小一点比如8或16这样即使显存不够也能跑。等确认没问题了再用tools/dist_train.sh启动多卡训练。多卡训练时有个坑学习率要按卡数线性缩放。比如单卡时学习率是0.014卡训练时应该设成0.04。这是因为多卡训练时每个卡上的batch_size不变但总的batch_size变成了4倍梯度更新的频率也变成了4倍。如果不缩放学习率训练会变得不稳定。还有一个坑是work_dir的冲突。多卡训练时每个进程都会往work_dir里写日志和权重如果路径没区分开会互相覆盖。mmclassification的启动脚本会自动给每个进程分配不同的临时目录但最终保存的权重是合并的这个不用太担心。4.3 训练过程中的监控与中断恢复训练启动之后你得盯着点。mmclassification默认用TensorboardVisBackend记录日志你可以用tensorboard --logdir work_dir在浏览器里看loss曲线和精度曲线。我一般关注三个指标loss是否稳定下降、accuracy是否稳步上升、lr是否符合预期。如果训练中途断了可以用--resume参数从最近的检查点恢复。但这里有个细节恢复训练时数据集的随机种子状态不会自动恢复。也就是说恢复后的数据顺序和中断前不一样。这对大多数场景没影响但如果你在做严格的对比实验最好固定随机种子或者接受这个微小的差异。还有一个实用技巧是设置CheckpointHook的save_best参数让训练过程中自动保存验证集上表现最好的权重。这样即使训练后期过拟合了你也能拿到中间某个时刻的最优模型。default_hooks dict( checkpointdict( typeCheckpointHook, interval1, save_bestaccuracy/top1, rulegreater, ), )5. 精度上不去的排查思路与实战经验自定义数据集训练时最让人头疼的不是报错而是loss正常降但验证精度死活上不去。这种情况通常不是代码bug而是数据或配置的细节出了问题。我总结了一套排查流程按这个顺序走基本能定位到原因。5.1 先确认数据本身有没有问题第一步永远是可视化检查。把训练pipeline去掉增强只保留LoadImageFromFile和PackInputs然后写个脚本把前几十张图存出来人眼过一遍。我遇到过标注文件和图片对不上的情况标注文件里写的是cat/001.jpg但实际图片在cats/001.jpg路径拼接后读到了另一张图模型学到的全是错的。第二步是检查类别分布。如果某个类别的样本数特别少模型会倾向于预测多数类少数类的精度会很低。这时候要么补充数据要么在损失函数里加类别权重。mmclassification的CrossEntropyLoss支持class_weight参数传一个和类别数等长的列表就行。第三步是确认标签范围。用numpy.unique统计一下所有标签看看是不是从0开始连续编号。如果标签是1到10但num_classes设成了10那标签10就会越界如果设成了11又多了一个空类模型会浪费容量去学一个不存在的类。5.2 学习率和batch size的匹配关系学习率和batch size是绑定的。业界有个经验公式学习率与batch size的平方根成正比。比如batch size从32增加到128学习率可以从0.01增加到0.02。但这个公式不是绝对的还得看优化器和数据集。我见过有人用batch size 8配学习率0.1结果训练初期loss直接飞了。也见过batch size 256还用0.001的学习率训练慢得像蜗牛。我的建议是先用官方配置的学习率跑一遍如果loss震荡就调小如果loss降得太慢就调大每次调整幅度不超过3倍。还有一个容易被忽略的点是warmup。mmclassification的配置文件里通常有param_scheduler里面会配一个线性warmup前几百个iteration学习率从0慢慢升到设定值。这个对稳定训练很有帮助尤其是batch size比较大的时候。如果你自定义数据集时把warmup去掉了训练初期可能会不稳定。5.3 预训练权重加载的坑用预训练权重能大幅加快收敛但加载的时候有几个坑。第一类别数不匹配。ImageNet预训练模型是1000类你的数据集可能只有10类加载时最后一层的权重形状对不上mmclassification会自动跳过不匹配的层但会打印警告。你得确认跳过的确实是分类头而不是其他关键层。第二归一化参数不匹配。预训练模型用的是ImageNet的均值和标准差如果你的数据集是灰度图或红外图直接套用ImageNet的归一化参数效果会很差。这时候要么自己统计数据集的均值和标准差要么把预训练权重的第一层改成单通道。第三加载方式。mmclassification支持load_from和pretrained两种方式。load_from是加载整个模型的权重pretrained是只加载backbone的权重。如果你只想用backbone的预训练特征用pretrained如果想在之前训练的基础上继续训练用load_from。# 只加载backbone预训练权重 model dict( backbonedict( typeResNet, depth18, pretrainedhttps://download.openmmlab.com/..., ), ) # 加载整个模型的权重继续训练 load_from work_dir/epoch_10.pth注意pretrained的URL如果填的是本地路径要确保路径存在且格式正确。我见过有人把相对路径填进去结果训练时找不到文件backbone随机初始化训练了半天精度还不如瞎猜。6. 把自定义数据集流程沉淀成可复用的模板每次做新项目都从头写Dataset和配置效率太低。我的做法是沉淀一套模板新项目来了改几个字段就能跑。模板的核心是把变化的部分参数化把不变的部分固化下来。6.1 模板目录结构的组织方式我习惯把模板组织成这样的结构my_template/ ├── configs/ │ ├── base_config.py │ └── my_dataset_config.py ├── datasets/ │ ├── __init__.py │ └── my_dataset.py ├── tools/ │ ├── train.py │ └── visualize.py └── data/ ├── train.txt └── val.txtbase_config.py里放通用的配置比如优化器、学习率调度、训练轮数。my_dataset_config.py继承base_config.py只覆盖数据集相关的字段。这样新项目来了复制一份my_dataset_config.py改改data_root、ann_file、classes就能用。my_dataset.py里放自定义的Dataset类如果项目用的是标准的txt格式直接继承CustomDataset就行不用自己写。visualize.py是我自己加的用来快速检查数据加载是否正确每次换数据集先跑一遍这个脚本能省很多调试时间。6.2 配置继承中的覆盖规则mmclassification的配置继承用的是_base_字段子配置里的同名字段会覆盖父配置。但这里有个细节字典是递归合并的列表是整体替换的。比如父配置里train_pipeline有5个变换子配置里只写了3个那最终生效的是子配置的3个而不是合并成8个。这个规则在改pipeline的时候特别容易踩坑。我见过有人想在父配置的pipeline基础上加一个变换结果子配置里只写了新增的那个变换把父配置的5个全丢了。正确的做法是把完整的pipeline都写出来或者用_delete_True先删掉再重写。# 正确做法完整写出pipeline train_pipeline [ dict(typeLoadImageFromFile), dict(typeRandomResizedCrop, scale224), dict(typeRandomFlip, prob0.5), dict(typePackInputs), ] # 或者用_delete_清空后重写 _base_ [../base_config.py] train_pipeline [ dict(typeLoadImageFromFile), dict(typeMyCustomTransform), dict(typePackInputs), ]6.3 从模板到新项目的迁移清单新项目来了按这个清单走一遍基本不会漏东西。确认数据格式是txt还是json图片路径怎么组织标签是单标签还是多标签准备标注文件确保路径和标签正确用脚本统计一下类别分布。修改配置文件改data_root、ann_file、classes、num_classes确认pipeline的输入尺寸和模型匹配。跑可视化脚本检查数据加载是否正确图片和标签是否对应。单卡小batch调试确认loss能降精度能涨。多卡正式训练调整学习率和batch size启动训练。评估与导出用测试集评估导出模型权重。这套流程我跑了不下二十个项目从几百张图的小数据集到几十万张图的大数据集都适用。关键是把每个环节的检查点固定下来不要跳步。跳步省下的时间后面排查问题的时候会加倍还回去。6.4 一个容易被忽略的收尾工作训练完了别急着关电脑还有几件事要做。第一在测试集上跑一遍最终评估确认精度符合预期。第二把配置文件和权重一起归档否则过两个月你根本想不起来这个权重是用什么配置训出来的。第三记录数据集的版本和预处理参数如果数据集更新了旧权重可能就不适用了。我习惯在work_dir里放一个README.md记录这次训练的关键信息数据集版本、配置文件路径、最终精度、训练时长、硬件环境。这个习惯帮我省了很多次“这个权重到底能不能用”的纠结。# 训练记录 - 数据集my_dataset_v2 - 配置文件configs/my_dataset_config.py - 最终精度top1 92.3% - 训练时长4小时4卡 - 硬件4 x GPU - 备注使用了随机旋转增强学习率0.04这套模板和流程不是一成不变的每做一个新项目我都会根据实际情况调整。但核心思路是不变的把重复的工作模板化把变化的工作参数化把关键的检查点固定化。做到这三点自定义数据集训练就不再是每次都要从头踩坑的苦差事了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑