ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

DINOv2工业缺陷检测落地实践:从自监督特征到小样本微调全流程

DINOv2工业缺陷检测落地实践:从自监督特征到小样本微调全流程 简介DINOv2是MetaAI基于Vision Transformer构建的自监督视觉模型在工业缺陷检测领域具有无需大量标注即可提取通用特征的优势。这套可运行源码包专为AI算法工程师与工业视觉研究者设计压缩包大小约13KB内含3个文件包括inscode运行配置、HTML说明文档以及.gitignore工程管理文件虽体量精简但覆盖了运行入口、说明文档与工程规范便于对照理解项目组织。模型采用对比学习策略通过同一图像不同变换构造正样本对在无标注数据上学习判别性表征特别适合表面缺陷、划痕、异物等小样本异常识别场景。目前已有99人学习该资源对于希望在低成本条件下验证自监督视觉模型、快速上手工业质检原型开发的开发者而言这份源码包能够有效缩短从理论到实践的路径同时其通用特征供后续分类、分割等下游任务迁移复用为算法迭代留出充分空间。 做工业视觉这么多年我越来越觉得多数检测项目卡住的不是算法模型本身而是“特征怎么提”和“缺陷怎么定义”这两件事。DINOv2 开源之后我花了大半个月把整套流程跑通整理成了一个带可运行源码的工业缺陷检测项目。这篇文章不是讲论文而是一份能直接抄作业的落地记录包括环境配置、数据准备、微调训练、推理可视化和一箩筐踩坑经验。如果你正被表面划痕、脏污、裂纹这类小样本缺陷折磨这篇内容或许能帮你少走不少弯路。1. 项目定位DINOv2 做工业缺陷检测的思路1.1 工业缺陷检测的难点工业在线质检本质上是一个“找异常”的问题。表面划痕、凹坑、脏污、边缘破损这些缺陷每个批次都不一样甚至同一片钢材上都没有完全相同的两条划痕。传统方案靠人工设计特征后来用 CNN 做目标检测和分割效果稳定了但遇到三个问题难解一是缺陷样本数量少正负样本极端不平衡二是缺陷形态多变一个类别内部差异大得离谱三是现场要求高召回误检又得压到极低。这批难题恰好是 DINOv2 这种自监督视觉模型最能补位的地方。DINOv2 是 Meta 在 2023 年开源的视觉特征学习模型核心是用自监督方式在大量无标注图片上预训练学到的特征具有非常强的语义和几何一致性。放到工业场景里它就像给模型提供了一个“见过世面”的底座哪怕下游只有几十张缺陷图也只需要少量微调就能把特征迁移过来效果往往比从头训练的 CNN 好一截。虽然 DINOv2 本身不解决检测框怎么画的问题但把它作为骨干网络接上检测头或分割头整个链路就顺了。1.2 为什么选 DINOv2 而不是纯自研 CNN直接回答一个经常被问的问题我用 YOLOv8 跑得好好的为什么要换 DINOv2我的判断标准很简单一是目标是否依赖纹理和局部上下文。划痕、裂纹、脏污这类靠局部纹理说话的缺陷DINOv2 的自监督特征相当敏感尤其是 Patch 级别的特征。二是样本是否稀缺。如果你的缺陷样本只有一两百张从头训 CNN 很容易过拟合用预训练特征再加一个轻量头收敛快且稳。三是是否需要随时加新类别。DINOv2 这类基础模型加新类别的成本低只需要在最后分类层做增量。当然它也有代价。模型体积比轻量 CNN 大推理速度要按现场设备评估。我在实际项目里通常用 DINOv2 做线下复杂质检场景的初筛模型或者做标注辅助产线实时级任务仍然会考虑蒸馏成轻量模型。也就是说DINOv2 更适合作为“高质量特征引擎”而不是最终打包进 PLC 之前的那个模型。1.3 源码仓库整体结构这个项目我整理成了可运行的源码仓库结构如下dinov2-defect/ ├── configs/ │ ├── train_config.yaml │ └── datasets/ ├── data/ │ ├── train/ │ │ ├── ok/ │ │ └── ng/ │ └── val/ │ ├── ok/ │ └── ng/ ├── src/ │ ├── dataset.py │ ├── model.py │ ├── trainer.py │ ├── infer.py │ └── utils.py ├── scripts/ │ ├── prepare_data.py │ └── export_onnx.py ├── pretrained/ │ └── dinov2_vitb14_pretrain.pth └── README.md核心不复杂就是“加载预训练权重 - 接检测头/分割头 - 微调训练 - 推理可视化”。数据目录按ok/ng划分训练脚本自动读取子文件夹作为二分类标签如果要做缺陷定位我额外预留了分割头配置。这套结构对第一次接触 DINOv2 的朋友来说比直接看官方仓库友好得多官方仓库更偏研究和特征提取演示我这个版本直接面向缺陷检测落地。2. 环境配置与依赖安装2.1 软硬件要求跑 DINOv2 微调核心不是算力多强而是显存和 torch 版本匹配。我测试过的稳定组合是Python 3.9 或 3.10PyTorch 1.13 或 2.xCUDA 11.7/12.1GPU 显存建议 10G 以上微调 ViT-B/14 时 batch size 4、分辨率 518x518 大约需要 7-8G如果显存只有 6G也能跑但需要梯度累积、混合精度和冻结部分参数。下面是推荐依赖清单conda create -n dinov2 python3.10 -y conda activate dinov2 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 pip install opencv-python pillow numpy tqdm pyyaml scikit-learn matplotlib需要说明的是PyTorch 2.0 以上版本对 DINOv2 的支持比较省心算子都集成在 torchvision 里了。预训练权重文件一般 300M 到 1G建议放在pretrained/目录并写进 gitignore避免每次同步仓库都重新下载。2.2 快速跑通 demo配好环境后最快的验证方式是直接跑我仓库里的 demo 脚本git clone https://github.com/yourname/dinov2-defect.git cd dinov2-defect python scripts/prepare_data.py --source ./examples --out ./data python src/train.py --config configs/train_config.yaml python src/infer.py --weights ./outputs/best.pth --img ./data/val/ng/001.png跑通这一步说明环境、权重、数据处理链路都正常。第一次跑我建议先用官方预训练权重直接做推理不做任何微调看一眼特征对缺陷的反应。后面再去调整参数这样能减少“环境背锅”的排查成本。说实话这比一上来就训练更容易发现配置问题。3. 数据处理与缺陷标注3.1 数据格式与目录组织工业缺陷检测数据格式决定了后续流程能走多远。我做项目的习惯是两边兼顾先按文件夹分类做二分类快速验证模型可用性再用 COCO 或 YOLO 格式做目标框微调。本项目里我同时支持两种模式分类模式data/train/ok,data/train/ng检测模式data/train/images/,data/train/labels.jsonCOCO如果你的产线还没有标注数据建议先采集正常样本 300 张、缺陷样本 100 张做初版。只需要标大类不需要细分缺陷类型DINOv2 的特征足够把大类差异拉得很开。后期要细分缺陷时再补标注模型底子还在重新训练的代价不会太大。3.2 预处理细节与数据增强DINOv2 对输入尺寸比较敏感。预训练时用的分辨率是 518x518我实测用 518 比 224 在细小划痕上效果好很多。224 会丢失太多纹理细节完全体现不出 DINOv2 的优势。所以默认训练分辨率我设成 518代价是显存差不多翻倍。增强方面要克制。工业质检不是日常图像分类过度旋转和色彩抖动反而会引入不真实的材质变化。我常用的增强组合是augmentation: resize: 518 random_crop: 0.2 hflip: true brightness: 0.1 contrast: 0.1 normalize: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225]注意不要加模糊和随机擦除。缺陷检测的任务就是找微小异常你把它擦掉或者模糊了模型就学偏了。这是我在一次手机中框表面检测项目里踩出来的教训用 RandomErasing 后召回率掉了近 8 个点。4. 模型微调与训练实现4.1 检测头设计从特征到缺陷决策DINOv2 起到骨干特征提取的作用真正决定检测效果的是下游头。我这里提供两种方案分类头用AdaptiveAvgPool - Linear - Softmax用于判断“有没有缺陷”分割头保留所有 Patch Token接一个FCN Head或U-Net风格的解码器输出像素级缺陷区域。实际项目里我通常会先用分类头跑通看图片级 AUC。如果 AUC 已经很高说明 DINOv2 特征足以支撑检测再升级到分割头定位缺陷。下面是一个简洁的分类模型实现import torch import torch.nn as nn from torchvision import models class DinoV2Classifier(nn.Module): def __init__(self, num_classes2, feature_dim768): super().__init__() # 用 torchvision 的 dinov2 权重也可以换成官方加载方式 self.backbone models.dinov2_vitb14(weightsmodels.DINOV2_ViTB14_WEIGHTS) self.head nn.Sequential( nn.AdaptiveAvgPool1d(1), nn.Flatten(), nn.Linear(feature_dim, num_classes) ) # 可选冻结 backbone 前几层 for param in self.backbone.parameters(): param.requires_grad False def forward(self, x): feat self.backbone(x) # (B, N1, D) cls feat[:, 0] # 取 [CLS] token return self.head(cls.unsqueeze(1))单纯分类头的结构很轻微调时通常只训练 headbackbone 冻结。但如果缺陷类型特别多或者缺陷长得很像正常纹理就需要解冻 backbone 的后 1/3 层用低学习率微调。4.2 训练参数与超参配置这里给出一个我测试过比较稳的训练配置model: name: dinov2_vitb14 num_classes: 2 pretrained: true freeze_backbone: true train: batch_size: 8 epochs: 30 lr: 3e-4 weight_decay: 1e-4 lr_scheduler: cosine warmup_epochs: 3 loss: cross_entropy mixed_precision: true grad_accumulation: 2 val_interval: 1如果你解冻 backbone学习率一定要降低我一般设置 backbone lr 为 head lr 的 0.1 倍例如backbone_lr3e-5。这个比例不是拍脑袋DINOv2 的预训练特征已经很成熟头部是随机初始化的若两者用同一学习率随机初始化头会有很大的梯度噪声容易反向扰乱预训练特征。4.3 损失函数与评价指标二分类场景直接CrossEntropyLoss但这要配合正负样本比例调整类别权重。缺陷样本少通常ok:ng是 3:1 甚至 5:1如果不加权重模型会偏向把图片判为 ok毕竟整体准确率也能到 80% 以上。我在训练代码里动态计算类别权重from sklearn.utils.class_weight import compute_class_weight labels [0, 0, 1] # 从数据集收集 weights compute_class_weight(balanced, classes[0, 1], ylabels) loss_fn nn.CrossEntropyLoss(weighttorch.tensor(weights, dtypetorch.float32).cuda())评价指标别只盯 accuracy。工业场景最重要的是 recall 和 precision 的平衡我习惯输出 Accuracy、Precision、Recall、F1、ROC-AUC。ROC-AUC 不依赖阈值能客观反映特征可分性。如果 AUC 在 0.95 以上后面只要调阈值就能满足现场要求如果 AUC 只有 0.8就要回去看数据或模型结构。5. 推理部署与可视化5.1 推理脚本用法训练完的模型我一般先导出 ONNX再拿到产线测试机上验证。导出的一个简单示例python scripts/export_onnx.py \ --weights outputs/best.pth \ --output outputs/best.onnx \ --img_size 518推理脚本支持两种输入模式单张图片和目录批量推理。单张推理示例python src/infer.py \ --weights outputs/best.pth \ --img data/val/ng/001.png \ --threshold 0.5 \ --save viz/输出结果包括预测类别、置信度和可视化图片。对于分类模型可视化会把缺陷区域用热力图高亮出来方便现场人员确认这台模型是“看”到了什么问题。5.2 热力图与 Grad-CAM 可视化很多人问DINOv2 不是做特征提取吗怎么知道它关注哪里我项目里实现了两种可视化Grad-CAM 对整个 ViT 输出层做梯度加权适合快速确认模型有没有看错区域Patch Similarity 利用 DINOv2 的 patch token 和 [CLS] token 的相似度画注意力图谱对纹理类缺陷更敏感。实际调试时如果一个样本被误判成 ok先看热力图是否集中在缺陷位置。如果热力图在背景上那说明特征是学到了但分类头没选对关键 patch如果热力图根本没亮说明模型可能过拟合了别的特征需要检查数据标注是否干净。6. 踩坑记录与常见问题排查6.1 显存不足与 OOM这是被问得最多的一个问题。DINOv2 ViT-B/14 在 518x518 输入下单卡 12G 训练 batch size 8 压力很大。解决方案优先级如下降低输入分辨率到 416 或 384先验证逻辑再上 518开启混合精度显存能省 30% 左右缩小 batch size 到 2配合grad_accumulation4保持等效 batch冻结 backbone只训练 head如果还不行用torch.utils.checkpoint对 ViT Encoder Block 做激活检查点。我实际用过torch.utils.checkpoint效果明显但会稍微拖慢训练速度。显存是硬指标省显存通常以时间和复杂度为代价。6.2 模型不收敛或过拟合不收敛的第一反应不是调学习率而是检查数据 pipelinenormalize参数是否和预训练一致DINOv2 官方是 ImageNet 归一化错一个 mean 和 std 都会导致训练震荡。标签是否对齐文件夹顺序乱会导致标签错位。样本是否严重不平衡类别权重没加会导致 loss 不降。过拟合则看验证 loss 是否在某个 epoch 后反弹。解决办法不是加大 dropout而是先检查训练集和验证集是否有同一张图片重复出现。工业场景经常因为数据没去重导致指标虚高我复现别人的项目时第一次就栽在这上面。6.3 DINOv2 权重加载失败常见的两个错误state_dict键名不匹配通常是因为把官方权重直接加载进 torchvision 模型但 torchvision 的键名前缀是backbone.开头官方是encoder.开头要写个映射函数转换。预训练权重文件损坏或版本过旧核对 SHA256 或者重新下载。我用的简单加载技巧是state_dict torch.load(pretrained/dinov2_vitb14_pretrain.pth) new_state_dict {} for k, v in state_dict.items(): new_key k.replace(encoder., backbone.) if k.startswith(encoder.) else k new_state_dict[new_key] v model.backbone.load_state_dict(new_state_dict, strictFalse)注意strictFalse只是临时排查用正式训练前还是要保证必要的键全部加载成功否则模型是在随机初始化下训练的。最后再分享一个我在实际项目中维持下来的习惯先构造一个 20 张的小数据集把训练、推理、导出全流程跑通再上全量数据。别小看这个习惯它能帮你过滤掉八成环境问题也能让你在老板面前快速拿出阶段结论。DINOv2 做工业缺陷检测最大的优势不是“模型更强”而是让团队从特征工程里解放出来把精力花在数据质量和现场部署上。等你们跑通这个链路大概率会有同感。本文还有配套的精品资源点击获取
返回列表