简介这份铁路轨道故障检测图像分类数据集面向轨道交通智能运维、工业视觉缺陷识别方向的算法学习者与工程实践者用于训练和评估轨道状态自动判别模型。数据已按正常与故障两类完成标注并预先划分训练集、验证集与测试集各类图片分目录存放便于直接接入主流深度学习框架开展分类实验。压缩包共803个文件以jpg与jpeg图像为主另含少量webp样本、一份标注json文件和一个可视化show脚本整体约278.41MB可运行脚本快速浏览样本分布与类别情况。目前已有212人学习下载。借助该数据集读者可完成从数据加载、类别可视化到CNN分类网络训练与调优的完整流程并结合作者主页的CNN改进、图像分割与YOLO目标检测项目进一步迁移到轨道缺陷检测等实际场景适合作为课程设计、毕业设计或工程验证的练手素材。1. 铁路轨道故障检测图像分类数据集800 张已标注样本能撑起什么铁路轨道故障检测这件事真正卡住大多数团队的往往不是模型结构而是数据。你手上可能有一堆巡检车拍回来的轨道图扣件缺失、弹条断裂、轨面裂纹、道床异物混在一起但没人标、没分类、没统一尺寸模型根本喂不进去。这个标题里的「约 800 张已标注图像分类数据集」本质是一份已经按类别分好文件夹、可以直接丢进训练脚本的小规模监督数据。它解决的是从零到一跑通 pipeline 的问题验证你的图像分类算法能不能在轨道场景下收敛验证数据标注规范是否合理验证部署链路是否通畅。适合两类人——一类是想快速搭一个铁路轨道故障检测 demo 的学生和初级工程师另一类是手里有大量原始巡检图、想先拿小样本试水再决定要不要大规模标注的工程团队。800 张不算多但分类任务里它足够跑通一次完整的训练、验证、推理闭环前提是你知道怎么用。2. 拆解这份轨道故障图像分类数据集类别、格式与标注质量判断拿到一份图像分类数据集第一件事不是急着写训练代码而是把它拆开看清楚。铁路轨道场景的图像分类和通用 ImageNet 那套差别很大背景高度重复都是道床、钢轨、扣件目标尺度小且细长类间差异可能只有几毫米的形变。所以判断这份数据能不能用得从类别定义、目录结构、标注一致性三个维度过一遍。2.1 轨道故障的典型类别划分与目录结构铁路轨道故障检测的常见分类体系一般围绕几个核心缺陷展开。扣件类故障包括扣件缺失、扣件松动、弹条断裂钢轨本体故障包括轨面裂纹、轨头剥离、轨腰锈蚀道床类包括道砟飞散、异物侵入。一份标注良好的分类数据集目录结构通常是按类别名建文件夹每个文件夹下放对应图像dataset/ ├── fastener_missing/ # 扣件缺失 │ ├── 0001.jpg │ ├── 0002.jpg │ └── ... ├── fastener_broken/ # 弹条断裂 ├── rail_crack/ # 轨面裂纹 ├── rail_corrosion/ # 钢轨锈蚀 ├── ballast_debris/ # 道床异物 └── normal/ # 正常轨道负样本这种ImageFolder式的结构是 PyTorch 生态里最通用的格式torchvision.datasets.ImageFolder可以直接读取不需要额外写解析代码。如果你的数据集不是这个结构比如所有图混在一个文件夹配一个 CSV 标签文件那就需要先转换。常见做法是写一个脚本按 CSV 把文件移动到对应类别目录。判断类别划分是否合理看两点一是类间是否互斥扣件缺失和弹条断裂不能是同一张图同时属于两类除非你做多标签但分类任务默认单标签二是类内是否一致如果「轨面裂纹」文件夹里混进了锈蚀图训练时模型会学到矛盾的特征。800 张数据如果分 5 到 6 类每类大概 130 到 160 张这个量级做迁移学习够用从零训练则偏少。2.2 标注质量自查三个能立刻跑的检查脚本标注质量决定了模型的天花板。图像分类的标注比目标检测简单——不需要画框只需要确认每张图归到正确的类别文件夹。但恰恰因为简单反而容易出问题文件名重复、图像损坏、类别不平衡、误标。下面三个检查脚本可以直接跑。第一个检查图像完整性和尺寸分布import os from PIL import Image from collections import defaultdict dataset_root ./dataset size_dist defaultdict(int) corrupt_files [] for cls in os.listdir(dataset_root): cls_dir os.path.join(dataset_root, cls) if not os.path.isdir(cls_dir): continue for fname in os.listdir(cls_dir): fpath os.path.join(cls_dir, fname) try: with Image.open(fpath) as img: img.verify() # 验证文件完整性 with Image.open(fpath) as img: size_dist[img.size] 1 # 统计尺寸分布 except Exception as e: corrupt_files.append((fpath, str(e))) print(尺寸分布:, dict(size_dist)) print(损坏文件:, corrupt_files)这段代码做两件事img.verify()检查文件是否能正常解码损坏的图会在训练时直接报错中断统计尺寸分布是为了决定后续 resize 策略如果尺寸差异极大比如 200x200 到 4000x3000 混在一起统一 resize 到 224x224 会损失很多细节需要考虑分块或裁剪。第二个检查类别数量和比例import os dataset_root ./dataset for cls in sorted(os.listdir(dataset_root)): cls_dir os.path.join(dataset_root, cls) if os.path.isdir(cls_dir): count len([f for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .png, .jpeg))]) print(f{cls}: {count} 张)如果某一类只有二三十张而其他类有一百多张这就是典型的长尾问题。800 张总量下类间比例超过 1:5 就需要处理否则模型会偏向多数类。处理方式后面章节会讲。第三个检查重复图像文件名不同但内容相同import hashlib import os dataset_root ./dataset hash_map {} duplicates [] for cls in os.listdir(dataset_root): cls_dir os.path.join(dataset_root, cls) if not os.path.isdir(cls_dir): continue for fname in os.listdir(cls_dir): fpath os.path.join(cls_dir, fname) with open(fpath, rb) as f: file_hash hashlib.md5(f.read()).hexdigest() if file_hash in hash_map: duplicates.append((fpath, hash_map[file_hash])) else: hash_map[file_hash] fpath print(f发现 {len(duplicates)} 组重复图像) for dup in duplicates[:10]: print(dup)重复图像在巡检数据里非常常见——同一段轨道连续拍摄多帧内容几乎一样。如果重复图分别落在训练集和验证集里验证指标会虚高这就是数据泄漏。发现重复后要么去重要么确保同一组重复图只出现在同一个 split 里。注意这三个脚本建议在划分训练集之前跑先清洗再划分顺序反了就得重来。3. 用迁移学习在 800 张轨道图上跑通分类从划分到训练800 张图从零训练一个 CNN 基本不可行参数太多、样本太少模型会直接记住训练集。迁移学习是标准解法拿 ImageNet 上预训练的权重做初始化只微调最后的分类层或者解冻部分层做小学习率微调。这一章给出完整的可复现流程。3.1 训练集/验证集/测试集划分与数据增强配置划分比例上800 张数据建议 7:1.5:1.5即训练集 560 张、验证集 120 张、测试集 120 张。如果某些类样本特别少用分层抽样保证每个 split 里类别比例一致。import os import shutil import random from sklearn.model_selection import train_test_split dataset_root ./dataset output_root ./split_dataset random.seed(42) for cls in os.listdir(dataset_root): cls_dir os.path.join(dataset_root, cls) if not os.path.isdir(cls_dir): continue files [f for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .png, .jpeg))] # 先分训练集和临时集 train_files, temp_files train_test_split( files, test_size0.3, random_state42) # 临时集再分验证和测试 val_files, test_files train_test_split( temp_files, test_size0.5, random_state42) for split_name, split_files in [(train, train_files), (val, val_files), (test, test_files)]: split_dir os.path.join(output_root, split_name, cls) os.makedirs(split_dir, exist_okTrue) for f in split_files: shutil.copy(os.path.join(cls_dir, f), os.path.join(split_dir, f)) print(划分完成)random_state42保证每次运行划分结果一致方便复现。train_test_split默认不打乱类别但因为我们是按类别循环处理的每个类别内部独立划分所以天然分层。数据增强方面轨道图像有几个特殊性图像可能是俯拍也可能是侧拍但翻转要谨慎——水平翻转对扣件方向敏感的场景可能产生不合理的样本。我一般用这些增强随机裁剪RandomResizedCrop、颜色抖动ColorJitter模拟不同光照条件、小角度旋转RandomRotation(10)。避免用垂直翻转和大幅旋转。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomRotation(10), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2, hue0.05), transforms.RandomHorizontalFlip(p0.3), # 低概率谨慎使用 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])Normalize用的均值方差是 ImageNet 的统计值因为骨干网络是在 ImageNet 上预训练的输入分布要匹配。RandomHorizontalFlip概率设到 0.3 而不是默认的 0.5是因为轨道扣件有方向性过度翻转会引入噪声。3.2 选 ResNet18 还是 EfficientNet800 张图的骨干网络取舍800 张图选骨干网络核心矛盾是容量和过拟合。ResNet18 参数量约 1100 万EfficientNet-B0 约 530 万MobileNetV3-Small 约 250 万。数据量小的时候参数少的模型反而更容易收敛到好的解。我的经验是800 张、5 到 6 类ResNet18 是安全选择EfficientNet-B0 在准确率上通常略好但训练更敏感需要更细致的学习率调度。如果最终要部署到边缘设备直接上 MobileNetV3 或 ShuffleNetV2牺牲一两个点准确率换推理速度。import torch import torch.nn as nn from torchvision import models def build_model(num_classes, backboneresnet18, pretrainedTrue): if backbone resnet18: model models.resnet18(pretrainedpretrained) # 冻结前几层只微调后面的 for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False model.fc nn.Linear(model.fc.in_features, num_classes) elif backbone efficientnet_b0: model models.efficientnet_b0(pretrainedpretrained) for name, param in model.named_parameters(): if features.7 not in name and features.8 not in name \ and classifier not in name: param.requires_grad False model.classifier[1] nn.Linear( model.classifier[1].in_features, num_classes) return model model build_model(num_classes6, backboneresnet18)冻结策略是关键layer4和fc参与训练前面的layer1到layer3冻结。这样做的好处是训练快、不容易过拟合因为底层特征边缘、纹理在 ImageNet 上已经学得很好了轨道图像同样需要这些基础特征。如果数据量涨到几千张可以解冻layer3一起微调。3.3 训练循环、学习率设置与早停策略训练配置上优化器用 AdamW学习率设 1e-3只训练 fc 层或 1e-4微调 layer4。batch size 在 800 张数据下设 16 或 32 都行显存够就 32。训练轮数不用设太多30 到 50 轮足够配合早停。import torch.optim as optim from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(./split_dataset/train, transformtrain_transform) val_dataset ImageFolder(./split_dataset/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) best_val_acc 0.0 patience 7 patience_counter 0 for epoch in range(50): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 验证 model.eval() correct, total 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total print(fEpoch {epoch1}, Loss: {running_loss/len(train_loader):.4f}, fVal Acc: {val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) patience_counter 0 else: patience_counter 1 if patience_counter patience: print(f早停于 Epoch {epoch1}) breakCosineAnnealingLR让学习率按余弦曲线下降比阶梯式下降更平滑在小数据集上效果通常更好。T_max30表示 30 轮完成一个周期。早停的patience7意味着验证准确率连续 7 轮不提升就停防止过拟合。filter(lambda p: p.requires_grad, model.parameters())这行很重要——只把需要更新的参数传给优化器冻结层的参数不参与更新省显存也避免意外修改。跑完之后加载best_model.pth在测试集上评估一次看混淆矩阵确认哪些类别容易混。轨道场景里扣件缺失和弹条断裂经常混因为两者视觉上都表现为「扣件区域异常」需要看局部细节才能区分。4. 轨道故障分类的避坑与排查标注、划分、训练里的真实翻车记录这一章记录的是我在轨道图像分类项目里实际踩过的坑每条按现象、原因、解决来写。800 张数据量小很多问题在大数据集上不明显但小数据集上会被放大。4.1 验证集准确率 95% 但上线就废数据泄漏的三种隐蔽形式现象训练时验证集准确率冲到 95% 以上但拿新拍的巡检图测试准确率掉到 60% 多。原因数据泄漏。三种常见形式——同一段轨道的连续帧被分到了训练集和验证集同一张图经过不同增强后分别进入两个 split文件名不同但内容完全相同的图分散在不同 split。轨道巡检数据尤其容易出第一种因为巡检车是连续拍摄的相邻帧差异极小。解决按「拍摄段」而不是按「单张图」划分数据集。如果数据里有拍摄时间戳或位置信息按时间段或里程段切分确保同一段的图只出现在一个 split 里。没有元信息的话用前面提到的 MD5 去重脚本先去掉完全重复的再用图像相似度如 SSIM检测近似重复。4.2 模型把所有样本都预测成多数类类别不平衡的处理顺序现象训练几个 epoch 后模型输出全部偏向样本最多的那一类少数类召回率接近零。原因类别不平衡。800 张图如果「正常轨道」有 400 张而「轨面裂纹」只有 50 张交叉熵损失会被多数类主导模型学到「全猜多数类」就能拿到不低的 loss。解决处理顺序很重要。第一步先确认不是标注错误导致的假不平衡比如把其他类误标成了正常类。第二步用加权交叉熵给少数类更高权重from sklearn.utils.class_weight import compute_class_weight import numpy as np labels [label for _, label in train_dataset.samples] class_weights compute_class_weight(balanced, classesnp.unique(labels), ylabels) class_weights torch.tensor(class_weights, dtypetorch.float).to(device) criterion nn.CrossEntropyLoss(weightclass_weights)第三步如果加权还不够用重采样——对少数类做过采样复制强增强或对多数类做欠采样。我一般优先加权因为不改变数据分布加权效果不够再考虑采样。4.3 增强过头反而掉点轨道图像不该用的三种增强现象加了数据增强后验证准确率比不加还低。原因增强方式和数据特性冲突。轨道图像里垂直翻转会产生物理上不存在的场景扣件不会倒挂大幅旋转会让钢轨的平行结构变得不平行CutMix 或 MixUp 在细粒度分类上可能把两个不同缺陷混成一张不合理的图。解决增强策略要保守。轨道图像我一般只用RandomResizedCrop、ColorJitter、小角度RandomRotation(10)水平翻转概率压到 0.3 以下。如果加了某类增强后验证指标下降先去掉它单独测一轮确认是不是增强的问题。4.4 训练 loss 正常但验证 loss 震荡batch size 和学习率的联动现象训练 loss 稳定下降但验证 loss 上下震荡准确率也跟着波动。原因batch size 太小比如 8加上学习率偏高1e-3每个 batch 的梯度噪声大模型在最优解附近来回跳。800 张数据用 batch size 8 的话一个 epoch 只有 70 个 step梯度更新次数多但每次方向不稳。解决两个方向调。一是增大 batch size 到 32梯度估计更稳定二是如果显存不够必须用小 batch把学习率降到 1e-4 并加 warmup。另一个容易忽略的点是num_workers设太小导致数据加载成为瓶颈GPU 等数据时训练节奏被打乱表现也是 loss 震荡。设num_workers4或更高配合pin_memoryTrue。4.5 测试集指标好但混淆矩阵暴露问题别只看准确率现象测试集整体准确率 88%看起来还行但混淆矩阵显示某一类召回率只有 40%。原因只看整体准确率会掩盖类别间的不平衡表现。如果多数类占比高整体准确率会被拉高少数类的差表现被掩盖。解决每次评估都打印混淆矩阵和每类的 precision/recall/F1。轨道故障检测里漏检把故障判成正常比误检把正常判成故障代价高得多所以少数类故障类的召回率比整体准确率更值得关注。如果某类召回率持续偏低优先检查该类标注是否有问题再考虑调整损失权重。from sklearn.metrics import classification_report, confusion_matrix model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.numpy()) print(confusion_matrix(all_labels, all_preds)) print(classification_report(all_labels, all_preds, target_namesval_dataset.classes))5. 把 800 张数据的价值榨干小样本下的进阶技巧与验证习惯800 张数据的天花板是明确的但通过几个技巧可以把它的价值推到接近上限。这一章讲的是我实际用过、确认有效的做法。第一个技巧是交叉验证替代单次划分。800 张数据做 5 折交叉验证每折训练集 640 张、验证集 160 张跑 5 次取平均指标。这样得到的评估结果比单次 7:1.5:1.5 划分可靠得多尤其在小数据集上单次划分的运气成分很大。代价是训练时间乘以 5但 800 张数据单次训练也就几分钟完全值得。from sklearn.model_selection import StratifiedKFold import numpy as np all_files [] all_labels_idx [] for idx, cls in enumerate(sorted(os.listdir(dataset_root))): cls_dir os.path.join(dataset_root, cls) if not os.path.isdir(cls_dir): continue for f in os.listdir(cls_dir): all_files.append(os.path.join(cls_dir, f)) all_labels_idx.append(idx) skf StratifiedKFold(n_splits5, shuffleTrue, random_state42) for fold, (train_idx, val_idx) in enumerate( skf.split(all_files, all_labels_idx)): print(fFold {fold1}: train{len(train_idx)}, val{len(val_idx)}) # 用 train_idx 和 val_idx 构建 Dataset 子集训练流程同上StratifiedKFold保证每折里类别比例和整体一致避免某一折里某个类完全缺失。第二个技巧是测试时增强TTA。推理时对同一张图做多次增强如原图、水平翻转、小角度旋转把多次预测的概率取平均。这个做法不增加训练成本通常能涨 1 到 2 个点在数据量小的时候效果更明显。def predict_with_tta(model, image_tensor, n_augments5): model.eval() probs torch.zeros(1, num_classes).to(device) with torch.no_grad(): # 原图 probs torch.softmax(model(image_tensor.unsqueeze(0).to(device)), 1) # 水平翻转 flipped torch.flip(image_tensor, dims[2]) probs torch.softmax(model(flipped.unsqueeze(0).to(device)), 1) # 小角度旋转 for angle in [-5, 5]: rotated transforms.functional.rotate(image_tensor, angle) probs torch.softmax( model(rotated.unsqueeze(0).to(device)), 1) return probs / (2 2) # 平均第三个技巧是错误分析驱动标注迭代。跑完模型后把验证集里预测错误的图单独挑出来看按错误类型归类是标注错了、图像质量太差、还是类间确实难分。如果发现某类错误集中说明该类样本不够或标注不一致下一轮标注优先补这类。800 张数据的第一版模型主要价值不是上线而是告诉你「还缺什么数据」。技巧预期收益额外成本适用条件5 折交叉验证评估更可靠训练时间 x5数据量 2000 张测试时增强1~2% 准确率推理时间 x3~5推理延迟不敏感错误分析迭代持续提升上限人工分析时间有标注迭代能力类别加权损失少数类召回 10~20%几乎无类间比例 1:3最后说一个习惯每次改完数据或训练配置固定跑一遍测试集并记录混淆矩阵不要只看准确率一个数。我吃过亏——有一次调了半天模型结构准确率涨了 2 个点结果看混淆矩阵发现是多数类涨了、少数类反而掉了上线后故障漏检率更高。从那以后我的评估脚本里混淆矩阵和每类召回率是必打印项准确率只是参考。800 张数据做铁路轨道故障检测模型结构的选择空间其实不大真正拉开差距的是数据清洗、划分策略和评估习惯。希望帮到你。本文还有配套的精品资源点击获取