简介这份铁路轨道故障检测图像分类数据集面向轨道交通智能运维、工业视觉缺陷识别方向的学习者与算法工程师用于训练和评估轨道状态自动判别模型。数据已按正常与故障两类完成标注并预先划分训练集、验证集与测试集同类图片集中存放便于直接接入CNN分类流程资源内附JSON标注文件与可视化脚本可快速核对类别分布与样本质量。压缩包共803个文件以779张jpg与20张jpeg图像为主体另含少量webp、1个json标注和1个py脚本整体约278.41MB规模适中适合课堂实验、课程设计或轻量级模型验证。目前已有212人学习下载。借助该数据集读者可省去采集与清洗成本直接开展分类网络训练、对比实验与改进验证并结合作者主页的CNN改进与目标检测内容延伸实践。1. 铁路轨道故障检测图像分类数据集约800张已标注图片二分类任务直接开跑铁路巡检这块视觉方案落地最卡脖子的从来不是模型结构而是没有一批干净、分好、能直接喂进 DataLoader 的图。我手上这份铁路轨道故障检测图像分类数据集约 800 张实拍图已经按正常、故障两类标注完毕训练集、验证集、测试集各自独立存放目录结构规整到可以直接ImageFolder读。它解决的就是从零攒数据这个最耗时的环节——你拿到手就能跑通一条 CNN 分类基线把精力放在网络改进和调参上而不是花两周去拍图、清洗、标注。适合做铁路巡检算法验证的工程师、拿它当课程设计或毕设数据源的学生以及想快速验证某个 backbone 在工业缺陷场景下表现的算法同学。类别数就 2 类具体以资源里的 json 文件为准别自己脑补第三类。2. 数据集结构与标注格式先看清目录再动手2.1 目录组织与划分逻辑这份数据最省心的地方在于划分已经做好了。常见做法是根目录下三个子文件夹train、val、test每个子文件夹里再按类别名分子目录形如dataset/ ├── train/ │ ├── normal/ │ └── fault/ ├── val/ │ ├── normal/ │ └── fault/ └── test/ ├── normal/ └── fault/这种结构是 PyTorchtorchvision.datasets.ImageFolder的默认约定——它会把每个子目录名当作类别标签按字母序编号。也就是说fault和normal谁编号为 0、谁为 1取决于字母顺序fault在前。这一点在算混淆矩阵、写类别映射时一定要确认否则你打印出来的类别 0 准确率可能跟你想的正好反了。json 文件里通常记录了类别名到索引的映射以及每类的样本数动手前先打开看一眼别跳过。划分比例上约 800 张的体量常见是 7:1.5:1.5 或 8:1:1。样本量不算大所以验证集和测试集可能各只有几十张评估指标的波动会比较明显。我的习惯是训练时固定随机种子评估阶段多跑几次取均值别拿单次结果下结论。2.2 标注格式与可视化脚本分类任务的标注其实就是文件夹归类没有边界框那种坐标文件所以不存在标注框偏移、漏标的问题但也意味着你无法从标注里看出故障的具体位置——这是分类和检测的本质区别。资源里带了一个 show 脚本用来可视化数据集通常是随机抽样若干张图拼成网格显示方便你快速确认图像有没有损坏、类别有没有放错。跑可视化之前先确认依赖一般就是matplotlib、PIL、numpy。如果脚本用的是cv2读图注意 OpenCV 默认 BGR 通道直接plt.imshow会颜色发蓝得先转 RGB。这类小坑不致命但第一次跑看到一片蓝紫色容易慌。提示先跑 show 脚本再训练。我见过太多人直接开训跑到一半 loss 不降回头才发现某类里有几张图是纯黑或者尺寸为 0 的坏图。3. 用 PyTorch 跑通二分类基线从 DataLoader 到训练循环3.1 构建 Dataset 与数据增强分类任务的标准起手式就是ImageFolder加transforms。铁路轨道图像的特点是背景相对固定道床、钢轨、扣件故障区域往往只占画面一小块所以增强策略要克制——过度随机裁剪可能把故障区域裁掉反而制造标签噪声。import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 训练集增强翻转轻微颜色抖动不做大角度旋转 train_tf transforms.Compose([ transforms.Resize((224, 224)), # 统一到骨干网络常用输入尺寸 transforms.RandomHorizontalFlip(p0.5), # 轨道左右翻转物理上合理 transforms.ColorJitter(brightness0.2, contrast0.2), # 应对不同光照 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # ImageNet 统计量 ]) # 验证/测试集只做确定性变换保证评估可复现 val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.ImageFolder(dataset/val, transformval_tf) test_ds datasets.ImageFolder(dataset/test, transformval_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) test_loader DataLoader(test_ds, batch_size32, shuffleFalse, num_workers4) print(train_ds.classes, train_ds.class_to_idx)逻辑说明训练集用RandomHorizontalFlip和ColorJitter做轻量增强验证和测试集保持确定性变换这样评估结果才有可比性。Normalize用的是 ImageNet 的均值和方差因为后面大概率要加载预训练权重输入分布得对齐。class_to_idx打印出来务必核对确认fault和normal的索引跟你预期一致。参数说明batch_size32对 800 张的体量比较合适太大梯度更新次数少太小 BN 层统计不稳。num_workers在 Windows 上如果报错就设成 0这是 DataLoader 多进程在 Windows 下的老问题。Resize到 224 是给 ResNet 系列用的如果你换别的骨干网络输入尺寸要跟着改。3.2 迁移学习训练循环与关键参数800 张图从头训一个 CNN 基本没戏过拟合会非常严重。标准做法是加载 ImageNet 预训练权重冻结浅层只微调后面的层或者干脆全网络小学习率微调。import torch.nn as nn import torch.optim as optim from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) # 加载预训练 ResNet18替换最后的全连接层为二分类 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 2) # 2 类输出 model model.to(device) criterion nn.CrossEntropyLoss() # 只优化需要梯度的参数学习率用 1e-3 起步 optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) for epoch in range(30): model.train() running_loss 0.0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 验证阶段 model.eval() correct, total 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) preds model(imgs).argmax(dim1) correct (preds labels).sum().item() total labels.size(0) acc correct / total print(fEpoch {epoch1}, Loss {running_loss/len(train_loader):.4f}, Val Acc {acc:.4f}) scheduler.step()逻辑说明resnet18换成二分类头CrossEntropyLoss内部自带 softmax所以模型输出直接是 logits不用手动加 softmax。优化器只更新requires_gradTrue的参数如果你后续想冻结浅层在训练前把对应层的requires_grad设为 False 即可。StepLR每 10 个 epoch 把学习率降为原来的 0.1帮助后期收敛。参数说明学习率 1e-3 是 Adam 微调的常用起点如果 loss 震荡厉害就降到 1e-4。epoch 数 30 对这个小数据集够用配合早停更好。batch_size和前面保持一致。评估时model.eval()和torch.no_grad()必须加否则 BN 层和 Dropout 行为不对显存也白占。3.3 类别不均衡与评估指标选择铁路轨道数据里故障样本通常少于正常样本如果直接看准确率模型全预测正常也能拿到很高的数这就翻车了。先统计一下每类数量from collections import Counter counts Counter([label for _, label in train_ds.samples]) print(counts) # 例如 Counter({0: 320, 1: 240})如果两类差距超过 2:1建议在CrossEntropyLoss里加weight参数或者用WeightedRandomSampler重采样。评估指标上别只看 accuracy把每类的 precision、recall、F1 都打出来故障类的 recall 尤其重要——漏检一个故障的代价远大于误报。注意测试集只在最后评估一次不要拿测试集调参否则你报出来的指标就是自欺欺人。4. 避坑与排查这份数据集上手时最容易栽的五个地方4.1 现象训练 loss 一直不降准确率卡在 50% 左右原因最常见的是标签索引搞反了或者ImageFolder读到的类别顺序跟你以为的不一样导致模型在学一个矛盾的映射。其次是图像本身有问题比如某些图全黑、尺寸异常。解决先打印train_ds.class_to_idx和Counter统计确认类别和数量。再跑 show 脚本肉眼过一遍图。如果都没问题把学习率降到 1e-4 重试Adam 在微调时 1e-3 偶尔会发散。4.2 现象验证集准确率远高于测试集原因验证集和测试集划分时可能分布不一致比如验证集里正常样本偏多或者测试集里有训练时没见过的拍摄条件。800 张的体量下这种波动很常见。解决检查三个子集的类别比例是否接近。如果差距大重新分层划分。评估时多跑几个随机种子取平均别拿单次结果当结论。4.3 现象DataLoader 在 Windows 上报BrokenPipeError或卡死原因num_workers 0时 Windows 的多进程机制跟 Linux 不同容易出问题。解决把num_workers设成 0或者把训练代码包在if __name__ __main__:里。这是 Windows 下的老毛病不是数据集的锅。4.4 现象显存不够batch_size 降到 8 还是 OOM原因输入尺寸设太大或者模型没放到 eval 模式做验证验证阶段还在算梯度。解决确认验证循环里有torch.no_grad()。输入尺寸从 224 降到 128 试试铁路轨道故障区域通常比较明显128 也能用。再不行换更小的骨干比如resnet18换mobilenet_v3_small。4.5 现象模型在训练集上 99%验证集只有 70%原因过拟合。800 张图对 CNN 来说太少尤其你如果解冻了全部层。解决冻结浅层只训分类头或者加Dropout、权重衰减。增强策略可以适当加强但别加随机裁剪——可能把故障区域裁没。早停也是必须的验证 loss 连续几个 epoch 不降就停。5. 进阶技巧用混淆矩阵和 Grad-CAM 验证模型到底学到了什么跑通基线只是第一步真正要确认模型看对了地方得靠可视化。混淆矩阵能告诉你两类之间有没有系统性混淆Grad-CAM 能告诉你模型做判断时关注图像的哪个区域——如果它盯着背景而不是钢轨表面那这个模型上线就是灾难。先算混淆矩阵from sklearn.metrics import confusion_matrix, classification_report import numpy as np model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in test_loader: imgs imgs.to(device) preds model(imgs).argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(confusion_matrix(all_labels, all_preds)) print(classification_report(all_labels, all_preds, target_namestest_ds.classes))classification_report会把每类的 precision、recall、F1 都列出来重点看故障类的 recall。如果 recall 偏低说明漏检多得回去调损失权重或者增强故障样本。Grad-CAM 这块pytorch-grad-cam库最省事from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.image import show_cam_on_image import numpy as np # 取 ResNet18 最后一个卷积层作为目标层 target_layers [model.layer4[-1]] cam GradCAM(modelmodel, target_layerstarget_layers) # 取一张测试图 img_tensor, label test_ds[0] input_tensor img_tensor.unsqueeze(0).to(device) grayscale_cam cam(input_tensorinput_tensor)[0] # 反归一化后叠加显示 img_np img_tensor.permute(1, 2, 0).numpy() img_np (img_np * [0.229, 0.224, 0.225] [0.485, 0.456, 0.406]).clip(0, 1) visualization show_cam_on_image(img_np.astype(np.float32), grayscale_cam, use_rgbTrue)逻辑说明target_layers指定要可视化的卷积层ResNet 一般选layer4的最后一层分辨率够用又保留语义信息。GradCAM会返回一张跟输入同尺寸的热力图值越大表示该区域对当前预测越重要。叠加到原图上你就能看到模型是盯着钢轨表面的裂纹/异物还是盯着道床背景。参数说明use_rgbTrue是因为我们反归一化后是 RGB 顺序如果前面用 OpenCV 读图得注意通道。热力图阈值可以调但默认的就行。我自己的习惯是每次训完一个分类模型先看混淆矩阵确认没有系统性偏向再抽 10 张故障图跑 Grad-CAM。如果热力图集中在背景或者图像边缘说明模型在走捷径这时候加多少数据增强都没用得回去检查数据里是不是有背景泄漏——比如所有故障图都在某个特定光照下拍的模型其实学的是光照而不是故障。从那以后我每次拿到新数据集都强制先跑一遍可视化再开训省得白跑几十个 epoch。希望这份数据和这套流程能帮你少走点弯路。本文还有配套的精品资源点击获取