DINOv2工业缺陷检测实战：小样本下的模型微调与部署

简介DINOv2是由MetaAI提出的自监督视觉基础模型基于Vision Transformer架构在海量无标注图像上训练提取的通用视觉特征可直接用于工业缺陷检测等下游任务无需微调即可使用显著降低对大量标注数据的依赖。此源码包面向工业视觉算法工程师、研究人员及具备一定Python基础的学习者解决了缺陷检测场景中标注样本稀缺、模型泛化能力不足的实际问题。压缩包共3个文件以inscode运行入口、HTML说明文档和gitignore配置为主整体仅13KB轻量易部署适合快速验证DINOv2在异常检测上的效果。已有99人学习下载。使用者可借此掌握基于自监督对比学习的检测流程理解环境部署要点与特征提取逻辑并能在源码基础上替换自己的工业图像数据开展实验为制造业质量控制提供高效、低成本的视觉检测方案。1. 为什么工业缺陷检测要引入DINOv2而不是从零训练这几年做工业视觉项目的朋友应该都有同感缺陷检测这个方向最卡的通常不是算法模型本身而是数据。我接触过的产线项目里良品样本动辄几万张但缺陷样本能凑出几百张就算不错了。尤其是某些偶发缺陷比如铸件内部气孔、薄膜表面划痕、PCB焊点虚焊现场能采集到的坏图屈指可数标完更是少得可怜。在这种情况下从头训练一个检测模型无论是分类网络还是目标检测网络效果都很难做上去过拟合几乎是必然的。后来我换了个思路能不能不从头训练而是用一个已经在海量自然图像上预训练好的通用特征提取器然后把缺陷检测当做下游任务来做这个思路其实不算新过去几年大家用ImageNet预训练的ResNet做迁移学习也是这么干的。但真正让我决定全面转向DINOv2的是一次实际对比测试。同一批钢材表面的小划痕数据集用ResNet50迁移学习和用DINOv2特征做分类头微调在只有几百张缺陷图的情况下DINOv2这边的F1分数高了将近8个点。这个差距不是调参能追回来的。DINOv2是Meta AI在2023年发布的自监督视觉模型它的核心思路是通过自蒸馏的方式让模型在没有人工标注的海量图像上学习到通用的视觉特征。这个模型在图像检索、语义分割、深度估计等任务上都表现出了非常强的特征表达能力而且它的特征对视角变化、光照变化、遮挡都有很好的鲁棒性。这些特性恰好是工业缺陷检测最需要的产线上的拍摄角度、光照条件、产品表面状态千变万化一个只在特定数据集上训练出来的模型换个工位可能就失灵了。我见过的很多团队在刚接触DINOv2时容易陷入一个误区觉得它是自监督模型是不是可以直接拿原始特征图做缺陷分割不用训练了实际上DINOv2的预训练权重虽然很强但它的输出特征是通用的不是专门为缺陷检测设计的。要做工业落地还是需要一个轻量级的分类头或者分割头把DINOv2提取到的特征映射到具体的缺陷类别上。好消息是这个头的参数量很小训练成本极低几百张样本就能训得不错。这也是我下面要展开讲的核心内容。2. 项目结构与环境配置一套能直接跑通的代码需要哪些部件2.1 完整项目的文件树与模块职责我先把这个可运行项目的整体结构摆出来让大家对规模有个直观感受。整个项目不复杂核心代码也就四五个文件但每个文件承担的职责非常明确。dinov2-defect-detection/ ├── configs/ │ └── defect_config.yaml # 模型参数、训练参数、数据路径配置 ├── data/ │ ├── train/ # 训练集按类别分文件夹 │ │ ├── normal/ │ │ └── defect/ │ └── val/ # 验证集结构同上 ├── src/ │ ├── dataset.py # 数据加载、增强、归一化 │ ├── model.py # DINOv2特征提取器 缺陷分类头 │ ├── train.py # 训练主脚本 │ ├── inference.py # 推理脚本支持单张图片和批量预测 │ └── utils.py # 日志、指标计算、模型保存等工具函数 ├── weights/ │ ├── dinov2_vitb14_pretrain.pth # 预训练权重需提前下载 │ └── best_model.pt # 微调后保存的最佳模型 ├── requirements.txt └── README.md这里我要重点说一下configs/defect_config.yaml这个文件。很多开源项目的配置文件写得极其复杂动不动几十个参数实际上大部分参数你根本不用动。我这个项目的配置尽量精简核心就这几项model: backbone_name: dinov2_vitb14 # 可选: dinov2_vits14, dinov2_vitb14, dinov2_vitl14 image_size: 224 freeze_backbone: true # 冻结DINOv2主干只训练分类头 use_lora: false # 是否在主干上挂LoRA适配器 data: data_root: ./data batch_size: 32 num_workers: 4 train: epochs: 30 lr: 0.001 weight_decay: 0.01 loss: focal # 可选: ce, focal inference: save_heatmap: true # 是否保存特征热力图 threshold: 0.5这些参数不是拍脑袋定的每一个选择背后都有实际考量。比如freeze_backbone我默认设为true因为我们的训练数据量太少如果整个模型一起做反向传播DINOv2的主干很容易被少数缺陷样本带偏灾难性遗忘会非常明显。而loss选focal而不是普通的交叉熵是因为缺陷样本和正常样本的比例往往严重不均衡focal loss能有效缓解这个问题。2.2 环境安装与依赖清单依赖这块我用的是最省事的方案直接写进requirements.txttorch2.0.0 torchvision0.15.0 opencv-python4.8.0 numpy1.24.0 timm0.9.2 pyyaml6.0 scikit-learn1.3.0 tqdm4.65.0 matplotlib3.7.0这里面最关键的坑在timm的版本。DINOv2官方代码在加载权重时依赖timm库但timm版本迭代很快API变动频繁。我实测下来0.9.2这个版本最稳定换成新版可能会报一些奇怪的参数名错误。如果你之前装过其他版本的timm建议先跑一下代码如果报错再根据提示降级。DINOv2的预训练权重需要从官方仓库下载。这里有四种规格的模型我直接给个参数对比表模型规格Patch大小参数量特征维度显存占用batch32, 224pxViT-S/14142100万384约4GBViT-B/14148600万768约8GBViT-L/14143亿1024约16GBViT-G/141411亿1536约32GB在实际产线项目中我用的最多的是ViT-B/14这个规格。ViT-S太小特征表达力在缺陷检测这种细粒度任务上还是有点吃力ViT-L和ViT-G虽然更强但对推理机器的显卡要求太高工业现场部署的工控机一般带不动。ViT-B在特征质量和推理速度之间是个不错的平衡点。3. 从数据准备到模型微调关键实现细节逐段拆解3.1 数据集格式与加载逻辑这个项目的数据集组织方式非常朴素就是按类别分文件夹。训练集和验证集各自独立每个类别一个子文件夹。对于缺陷检测来说最少需要两个类别normal和defect。如果你的项目涉及多种缺陷类型比如划痕、气泡、脏污那就建立多个子文件夹。文件命名没有任何限制只要后缀是.jpg、.jpeg或.png就行。数据集加载时会根据文件夹名自动生成类别标签所以文件夹命名务必规范不要带空格或特殊字符。dataset.py里有几个让我印象深刻的处理细节。第一个是图像的归一化参数DINOv2预训练时用的均值和标准差是(0.485, 0.456, 0.406)和(0.229, 0.224, 0.225)和ImageNet保持一致。这个必须用对否则提取出的特征会受到很大影响。第二个是数据增强策略我采用的是随机水平翻转、随机垂直翻转、随机旋转±15度和随机亮度对比度调整。这里要特别提醒对于缺陷检测场景不要用太强的几何增强比如随机裁剪就不建议因为缺陷特征本身可能非常细微裁剪比例不合适容易丢失关键信息。3.2 模型定义DINOv2特征提取器 轻量分类头model.py是整个项目的核心模型定义非常简洁import torch import torch.nn as nn import timm class DINOv2DefectClassifier(nn.Module): def __init__(self, backbone_namedinov2_vitb14, num_classes2, freeze_backboneTrue): super().__init__() # 加载DINOv2预训练模型只保留特征提取部分 self.backbone timm.create_model( backbone_name, pretrainedFalse, num_classes0, # 去掉分类头 img_size224 ) # 加载官方预训练权重 state_dict torch.load(fweights/{backbone_name}_pretrain.pth) self.backbone.load_state_dict(state_dict) for param in self.backbone.parameters(): param.requires_grad not freeze_backbone # 获取特征维度 feat_dim self.backbone.embed_dim # 轻量分类头 self.classifier nn.Sequential( nn.Linear(feat_dim, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(256, num_classes) ) def forward(self, x): # DINOv2输出的是patch token序列 [B, N1, D] # 取CLS token作为全局特征 features self.backbone.forward_features(x) # features shape: [B, C] 或 [B, N, C]取决于模型输出格式 if features.dim() 3: features features[:, 0, :] # 取CLS token return self.classifier(features)这里有一个容易踩坑的点不同版本的timm对DINOv2的forward输出格式定义不太一样。有的版本forward_features返回的是[B, N, D]的序列有的版本已经帮你把CLS token提取好了返回[B, D]。为了避免出错代码里做了一个维度判断统一处理。这个细节看起来不起眼但实际操作中很多人就是卡在这里报维度不匹配的错误。分类头我用的是两层MLP中间加了ReLU和Dropout。可能有人会问为什么不直接用单层线性分类器其实对于简单的二分类正常vs缺陷单层线性分类器也够用。但如果缺陷类型比较多或者某些缺陷特征比较接近比如不同材质的划痕两层MLP的判别能力会更强一些。3.3 损失函数与训练参数的设计逻辑损失函数的选择是我在这个项目里花时间最多的地方之一。工业缺陷检测有个很现实的问题缺陷样本的数量可能只占总样本的10%甚至更少。如果直接用交叉熵损失模型会倾向于把所有样本都预测为正常类因为这样就能获得99%以上的准确率。为了解决这个问题我默认使用了Focal Loss。它的核心思想是让模型更关注那些难以分类的样本。从公式上看Focal Loss在标准交叉熵的基础上增加了一个调制因子(1-p)^γ其中p是模型对正确类别的预测概率γ是聚焦参数。当p接近1时样本容易分类调制因子趋近于0当p较小时样本难以分类调制因子较大损失贡献也就更大。在我的实践中γ设为2效果最好。另外配合使用类别权重给缺陷类一个更高的权重通常设为1 (正常样本数 / 缺陷样本数)。我遇到过的一个场景是正常样本2000张、缺陷样本180张这时候缺陷类的权重设为1 2000/180 ≈ 12。这个值看起来很大但配合相对较大的γ实际效果是稳定收敛的没有出现训练震荡。训练优化器我选的是AdamW学习率设0.001配合余弦退火调度器。因为在冻结主干的情况下只需要训练分类头那几百个参数学习率没必要用很低。如果后续你选择解冻主干做全量微调学习率必须降到0.0001以下否则主干特征会被破坏。训练过程通过早停法来防止过拟合验证集准确率连续5个epoch没有提升就停止训练。实测下来大多数项目在15到20个epoch时就能收敛到最优状态30个epoch的上限是足够的。4. 推理与可视化如何让模型输出在产线上真正可用4.1 推理流程与阈值选择训练完成后的推理流程相当直接。inference.py里实现了一个标准的推理管线读取图像、预处理、DINOv2特征提取、分类头预测、输出结果。这里有一个关键参数是分类阈值threshold。默认是0.5但实际使用时我建议不要直接用默认值。面对类别不平衡数据训练的模型预测为缺陷类的概率通常会偏保守也就是说正常样本的预测概率会非常接近1而缺陷样本的预测概率可能只在0.6到0.9之间。如果你把阈值设为0.5会有相当一部分缺陷被漏掉。我的做法是在验证集上计算不同阈值下的F1分数画出P-R曲线选取使F1分数最大化的阈值作为实际部署的阈值。这个操作听起来复杂但用sklearn的precision_recall_curve几行代码就能搞定。4.2 特征热力图把模型判断依据可视化save_heatmap参数控制一个很有价值的可视化功能特征热力图。这个功能不是锦上添花在产线调试阶段几乎是必须的。当你面对一个误检案例时光看模型的输出概率远远不够你得知道模型是根据图像哪个区域做出的判断。如果模型是因为背景纹理误判了缺陷你看到热力图就能立刻发现。实现原理不复杂。DINOv2输出的patch token包含了每个图像块的特征我们可以计算每个patch的特征向量与缺陷类别特征向量的相似度然后把这些相似度根据patch的位置映射回原始图像尺寸最后叠加到原图上得到热力图。这样就能直观地看到模型在关注哪些区域。我截取一段关键代码def generate_heatmap(self, image, target_class1): # 获取patch token特征 with torch.no_grad(): features self.backbone.forward_features(image.unsqueeze(0)) # features: [1, N1, D] patch_features features[0, 1:, :] # 去掉CLS token # 分类头的第一层权重可以视为特征重要性参考 head_weight self.classifier[0].weight[target_class] # [D] # 计算每个patch的响应分数 similarity patch_features head_weight # [N] # 还原为2D网格 grid_size int(patch_features.shape[0] ** 0.5) heatmap similarity.reshape(grid_size, grid_size) # 上采样到原图尺寸 heatmap F.interpolate( heatmap.unsqueeze(0).unsqueeze(0), size(self.image_size, self.image_size), modebilinear, align_cornersFalse ).squeeze() return heatmap在几次产线调试中这个热力图帮我快速定位了两个问题。一次是模型把产品表面的水渍误判成缺陷看了热力图发现模型关注的是图像角落的反光区域——那是拍摄光源的镜面反射产生的。另一次是模型对某类轻微划痕识别不稳定热力图显示模型同时关注了划痕和旁边的字符印刷区域导致特征被稀释。这些信息不通过热力图很难发现。4.3 评估指标不要只看准确率工业缺陷检测的评估指标必须谨慎选择。我见过不少项目在汇报时说准确率99.5%但实际效果一塌糊涂——因为缺陷样本只占0.5%模型把所有样本都预测为正常准确率也接近99.5%。这种虚高指标对产线完全没有参考价值。我建议每个项目至少报告四个指标精确率Precision、召回率Recall、F1分数和混淆矩阵。召回率尤其关键因为在缺陷检测场景中漏检比误检的代价大得多。一个漏检的缺陷产品流入下游可能造成批量退货而误检顶多多一次人工复检。我们在实际项目中的经验是对于影响产品质量的严重缺陷召回率最好能到98%以上精确率可以适当放宽到90%而对于不影响功能的表面瑕疵精确率和召回率平衡在95%左右即可。这个标准要根据具体工艺要求来定但无论如何用单一准确率指标考核模型是非常危险的。5. 踩坑记录与生产环境落地建议5.1 三个典型坑的完整排查链路坑一加载预训练权重时state_dict的key不匹配。第一次跑通这个项目的朋友十有八九会碰到这个问题。现象是load_state_dict报错提示Missing key(s): ... unexpected key(s): ...。排查思路是这样的先打印一下预训练权重的key集合再看一下timm模型内部的key集合对比差异。我遇到的情况是官方权重里有一些mask_token和相关的归一化层而timm模型的架构里没有这些键需要设置strictFalse跳过缺失的键。代码里我已经处理了这个情况但如果你在别的项目里复用DINOv2的权重这个坑大概率还会遇到。坑二训练时显存不足但batch size已经很小了。这通常是因为图像尺寸过大。DINOv2的ViT结构计算复杂度与图像分辨率成正比我遇到过一个客户要求直接检测1920x1080的原始图像batch size调到1也会OOM。解决思路是先缩放再检测比如缩放到448x448而不是直接硬扛。不过要注意缩放会丢失一定的细节对于特别细微的缺陷需要谨慎评估。另一种思路是用滑动窗口把原图切成多个patch分别检测最后合并结果。坑三验证集准确率很高但现场连续误报。这是最让人头疼的坑。排查链路一般是先检查现场图像和训练集图像的光照差异再用热力图看模型关注的区域最后检查是不是数据集划分出了问题比如同一批产品既进了训练集又进了验证集导致过拟合评估失实。我遇到过一次很隐蔽的问题训练集是用实验台拍摄的现场是产线传送带上的图像背景差异很大。DINOv2提取的特征确实对光照有一定的鲁棒性但背景纹理的差异仍然会带来明显的性能下降。解决方案是在现场采集一部分图像补充训练集哪怕只有几百张也能显著改善。5.2 从实验到产线部署的几条经验从实验环境到产线部署有几个容易被忽视但极其重要的细节。第一个是推理速度的验证时机。在Python环境下单张推理耗时可能还行但产线节拍是按秒甚至毫秒算的。建议在项目初期就用torchscript或ONNX导出模型做一次速度验证不要等到部署阶段再发现性能不达标。DINOv2这种ViT结构的模型导出ONNX会遇到一些兼容性问题主要是Transformer里的reshape和transpose操作在不同版本的ONNX Runtime里实现不一致需要多试几个版本。第二个是模型的版本管理。工业项目往往要长期运行模型迭代周期可能以月为单位。建议在保存模型时把训练数据、代码版本、模型参数全部打上标签录成一个清单。这一步看似简单但能省去后面无数次的溯源排查。第三个是对模型置信度飘移的监控。产线环境变化比如换了光源、换了相机会导致模型输出分布逐渐偏移。一个低成本的做法是每生产1000个产品统计一下模型输出的置信度分布如果整体偏移超过10%就说明图像分布变了需要重新校准甚至重新训练。这个监控脚本非常轻量但价值巨大能在不良品流出之前提前预警。5.3 后续扩展方向从分类到分割当前的代码框架做的是图像级别的二分类——整张图是正常还是缺陷。但在实际应用中很多场景不仅要知道有没有缺陷还要知道缺陷在哪里、有多大。这时候就需要把分类头换成分割头做像素级的缺陷检测。DINOv2其实很适合做这个扩展因为它的patch级特征天然保留了空间信息。你可以把模型的输出从取CLS token改为取全部patch token然后在上面接一个轻量的分割头比如FCN头或简单的上采样卷积对每个patch做二分类有缺陷/无缺陷最后组合成完整的缺陷分割图。这个扩展在代码层面改动不大但需要标注像素级的数据集人工成本会高不少。另外use_lora参数预留了LoRA适配器的扩展能力。如果你想在保持DINOv2主干权重不变的前提下让模型更好地适配特定产线数据可以在主干上挂几个LoRA低秩适配层只训练适配层参数。这样既保留了主干特征的通用性又能让模型更贴合目标场景。这个方向我在几个项目里做过实验效果比单纯微调分类头要好尤其是当目标场景与预训练数据分布差异较大时提升更明显。关于这个项目我个人的实际体会是DINOv2真正解决了工业视觉里小样本学习的痛点但它的价值要在一个完整的工程框架中才能释放。预训练权重只是起点你需要想清楚特征怎么接、损失怎么设计、阈值怎么调、错误怎么排查而所有这些单靠一份源码是远远不够的。希望这篇文章能帮你少走一些弯路。如果你在跑通源码的过程中遇到其他问题欢迎在评论区把报错信息贴出来我们一起排查。本文还有配套的精品资源点击获取