优化深度学习模型以在多变环境下检测列车车厢间隙

背景 #

我们团队从不同机位采集移动列车视频，拍摄背景、与轨道距离各不相同，核心任务是检测车厢间隙以收集车厢信息：

• 最初基于2000余张标注图+无间隙未标注图，用带默认数据增强的Yolov5训练检测模型，存在假阳性率高、低光环境性能差的问题；
• 后处理采用DBSCAN算法对含“车钩”的帧分组，通过平均置信度与标准差过滤低置信样本；
• 近期采集5万张多场景图片（含/无车钩），用现有模型动态标注：车钩置信度≥60%标记为GAP，<60%舍弃，无车钩标记为NO_GAP，基于此用Yolov8训练[GAP, NO_GAP]二分类器，但因NO_GAP涵盖场景过于繁杂，对模型泛化能力存疑；
• 目前考虑半监督学习、对比学习等技术，以及ViT（带分块机制）等架构，但相关经验不足，特提出以下问题：

问题1：推荐哪些深度学习架构或技术，提升多变光照与环境下的检测精度？ #

• 针对低光环境：
  • 数据层面做针对性增强：加入随机亮度调整、伽马校正、模拟低光噪点的增强策略，或用Retinex、CLAHE这类传统图像增强算法预处理低光样本后再喂给模型；
  • 模型层面选光照鲁棒性更强的架构：比如YOLOv8系列里的大模型（如YOLOv8l），或把YOLO的backbone替换为在低光数据集上预训练的ResNet/EfficientNet；
• 针对假阳性与复杂背景：
  • 优化数据增强：加入场景相关的定制增强，比如随机遮挡、背景替换（用不同轨道/环境的背景合成样本）、升级Mosaic增强为列车场景专属版本；
  • 半监督学习：用少量标注数据+大量未标注数据训练，比如FixMatch、MeanTeacher这类算法，让模型在未标注数据上学习通用特征，减少假阳性；
  • 后处理优化：除DBSCAN外，加入时序过滤——利用列车移动时车厢间隙位置的连续性，用卡尔曼滤波跟踪车钩位置，过滤帧间位置突变的假阳性；

问题2：是否适合继续采用分类任务，改用ViT（带分块机制）这类架构？有无可参考的具体实现案例？ #

• 是否适合：
  若核心需求是判断单帧是否存在车厢间隙（而非定位间隙位置），分类任务完全可行，但需解决NO_GAP类别的样本多样性问题——先对NO_GAP样本聚类，分成“无车钩的列车车身”“轨道背景”“其他干扰物”等子类，平衡各子类样本量后再训练，避免模型学到“非列车”特征而非“无间隙”特征；
  ViT擅长捕捉全局特征，若先裁剪出列车区域再输入，能有效捕捉车厢结构信息，适合这类任务；若用整图输入，易被背景干扰，建议先做目标检测定位列车区域，再对裁剪区域做二分类；
• 具体实现参考：
  • 用PyTorch实现ViT二分类：将输入图像resize至224x224，调用torchvision.models.vit_b_16作为backbone，替换最后的分类头为二分类层；
  • 铁路场景ViT落地思路：先过滤模糊、低质样本，用对比学习预训练ViT的铁路场景特征，再微调二分类头；
  • 轻量化替代方案：尝试YOLOv8的分类版本（YOLOv8-cls），工业场景泛化性已验证，比纯ViT更轻量化，训练推理速度更快；

问题3：大量未标注数据是否值得用自监督学习做预训练？在数据比例、算力、算法选择、停止时机上有哪些经验法则？ #

• 是否值得：非常值得。当前标注数据仅2000+，未标注数据达5万+，自监督预训练能让模型先学习铁路场景的通用视觉特征，再微调检测/分类任务，可大幅提升泛化能力；
• 经验法则：
  • 数据比例：未标注数据量至少是标注数据的5倍以上（当前25倍完全符合），同场景未标注数据越多，预训练效果越好；
  • 算力需求：若用ViT做自监督预训练，建议至少8张A100（或等效算力），训练周期1-2周；若用轻量化模型如EfficientNet，4张V100即可完成，训练周期3-5天；算力有限时优先选MoCo v2、SimCLR v2这类低算力需求的自监督算法；
  • 算法选择：优先选对比学习类算法，MoCo v2适合小算力场景，SimCLR v2适合大算力场景；若针对检测任务，可选用DetCo这类面向检测的自监督预训练算法；
  • 预训练停止时机：监控预训练损失曲线，当损失连续3-5个epoch不再下降时停止；或每隔几个epoch用预训练backbone微调下游任务，下游任务精度不再提升时停止预训练；

内容的提问来源于stack exchange，提问作者Bartłomiej Gładys