尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

DINOv2实战:用自监督预训练模型轻松搞定自定义图像分类

DINOv2实战:用自监督预训练模型轻松搞定自定义图像分类 做计算机视觉的同行应该都有这种体会数据标注太贵训练样本太少每次接到一个新的图像分类需求第一反应就是去翻ImageNet预训练权重。以前我常用ResNet50的预训练模型然后微调最后几层效果凑合但总觉得上限不高。后来换上了DINOv2情况有了质的改变。这篇文章记录的是我在自定义数据集上用DINOv2做图片分类预测的完整过程包括环境准备、数据组织、特征提取、分类器训练以及那些常规文档里不会写的坑。适合正在做图像分类、少样本识别、图像检索的工程师参考也适合刚接触自监督视觉模型、想找个靠谱方案入门的同学。DINOv2是Meta AI在2023年发布的自监督视觉模型系列。所谓自监督简单说就是模型在预训练阶段根本不需要人工标注的类别标签只靠图片本身就能学到通用的视觉特征。这听着有点像玄学但实际效果非常能打它学出来的特征向量可以直接拿来训一个简单的线性分类器在很多下游任务上就能超过传统“预训练微调”的路线。最让我喜欢的一点是它提取的特征是通用的也就是说你在自己的数据集上不需要大规模重新训练只要把图片过一遍模型得到特征向量再训练一个轻量分类头就够了。我这次拿一个二分类的猫狗数据集来演示整个流程。选这个数据集没有别的原因就是大家熟悉、复现方便把代码里的类别数改一下就适用于任何自定义任务。本质上的方法论是通用的你换成工业零件缺陷分类、遥感图像场景分类、医学影像筛查核心流程完全一样。1. DINOv2核心原理与选择理由1.1 没有标签DINOv2是怎么学会特征的DINOv2的前身是DINO和iBOT核心思路是做自蒸馏。你可以把它理解成让一个“老师模型”和一个“学生模型”同时看同一张图片的两个不同视角也就是经过不同裁剪、颜色扰动后的版本然后让学生模型的输出尽量去匹配老师模型的输出。老师模型的参数不是单独训练的而是对学生模型做指数移动平均得来的。这套机制很像一个学生不断去模仿一个“更好的自己”在模仿过程中模型被迫学会抓住图片中最本质、最稳定的内容而不是那些容易被扰动干扰的表面信息。在ViTVision Transformer架构下图片会被切成一个个固定大小的patch比如224x224的图按14x14的patch切就能得到16x16共256个patch。每个patch映射成一个token前面再加一个特殊的class tokenCLS token。训练完之后这个CLS token对应的输出向量就可以当作整张图片的全局特征表示。DINOv2的特征之所以通用性强是因为它在海量无标注图片上见过太多视觉模式学出来的特征层级非常丰富从边缘纹理到物体部件再到整体语义都有覆盖。1.2 相比传统迁移学习DINOv2赢在哪里传统做法是在ImageNet标注数据上训练一个分类模型然后拿到自己的任务上微调。这套路本身没问题但有两个隐性痛点一是学到的特征高度跟ImageNet的1000个类别绑定迁移到跟你数据分布差异很大的场景时需要微调的参数量很大二是如果自己的数据集很小微调很容易过拟合。DINOv2因为训练时不依赖标签学到的特征更偏向通用的视觉结构而不是特定类别的表面模式所以迁移到新任务时往往只需要一个线性分类头就能出很好的效果。方案需要标注数据自定义数据集上所需训练量特征通用性上手复杂度ResNet ImageNet微调需要需要训练整个网络或多层中等中等CLIP零样本/少样本需要文本描述基本不需要训练强但有文本依赖低DINOv2 线性探针需要少量标注只训练一个线性层极强低CLIP也是个很强悍的模型但它依赖图文配对数据如果你的任务类别很难用一两句话描述清楚或者你的图片风格跟CLIP训练数据差得很远效果反而不好。DINOv2只需要图片本身在纯视觉任务上更纯粹。这也是我在这个项目里选择DINOv2的直接原因。1.3 四个规格怎么选DINOv2官方提供四个规格的模型分别是ViT-S/14、ViT-B/14、ViT-L/14和ViT-g/14。字母后面的14表示patch size是14像素。从使用体验来说小模型速度快、显存占用少大模型精度上限更高但资源消耗也上去了。我建议普通用户先从ViT-B/14开始它在精度和开销之间最均衡。如果你的数据量不大、特征维度要求不高ViT-S/14完全够用。显存充裕且追求极致效果再上ViT-L/14。ViT-g/14参数量超过10亿一般个人项目没必要碰。2. 环境准备与自定义数据集组织2.1 软硬件环境这次实验我用的是一张RTX 309024GB显存但其实DINOv2 ViT-B/14的显存占用并不会让你有压力提取特征时即使把batch size设到64显存占用也就几个GB。如果你是8GB显存的卡把batch size调到16或者32也完全跑得动。CPU推理我也试过一张224x224的图大概要几百毫秒到一秒多做离线特征提取勉强能忍但实时预测就不建议了。软件方面Python 3.9以上PyTorch 2.xtorchvision再加上scikit-learn、matplotlib、Pillow这几个常用库就行。DINOv2的模型是通过torch.hub加载的所以需要能访问GitHub和下载权重文件。首次加载会自动下载预训练权重权重文件比较大ViT-B/14大概300多MB建议提前找个网络好的时间点装好后面实验会顺畅很多。安装依赖的命令很简单我直接用requirements.txt管理torch2.0.0 torchvision0.15.0 scikit-learn1.2.0 matplotlib3.7.0 pillow9.5.0 tqdm4.65.0 numpy1.24.02.2 数据集目录结构自定义数据集的组织方式非常自由DINOv2本身不关心你图片存哪只关心你喂给它的张量长什么样。但为了让代码逻辑清晰我建议用ImageFolder风格的组织方式data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ ├── cat_002.jpg │ │ └── ... │ └── dog/ │ ├── dog_001.jpg │ ├── dog_002.jpg │ └── ... └── val/ ├── cat/ │ └── ... └── dog/ └── ...如果你自己的数据是一堆乱放在文件夹里的图片但有一份CSV标注文件比如文件名和类别两列也是可以的。后面的Dataset类只需要按你的标注文件去读取就行。我这次演示先按目录结构来因为torchvision的ImageFolder直接支持这种格式代码最省事。我自己手工收集数据的时候习惯写个小脚本按类别把图片移动到对应目录顺便检查一下有没有损坏的文件否则训练到一半发现某张图解码失败很浪费感情。2.3 图像预处理细节DINOv2在预训练时输入图像会被缩放到224x224并且做了标准化的ImageNet均值和方差。这个预处理跟你自己的数据是什么风格关系不大建议直接照搬from torchvision import transforms transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])一个常见的疑惑是我的图片是灰度图怎么办最简单的办法是把单通道复制成三通道。更好的做法是先把灰度图做适当对比度增强或伪彩色映射但这些属于数据本身的预处理策略跟DINOv2无关。需要注意的是DINOv2的patch size是14224刚好是14的整数倍这是设计好的分辨率不建议随意改成其他尺寸。如果你强行用更大的图比如448x448虽然也能跑但位置编码是预训练时固定好的对不齐的位置编码会严重影响效果。所以除非你打算重新训练位置编码否则老老实实用224x224。2.4 自定义Dataset类写一个Dataset类其实不复杂核心就是返回图像张量和它的标签。如果直接用ImageFolder代码更短from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader train_dataset ImageFolder(data/train, transformtransform) val_dataset ImageFolder(data/val, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleFalse, num_workers4) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4)这里要注意shuffleFalse是因为我们要按顺序提取特征顺序在后面跟标签做对齐时非常关键。一旦shuffle了特征跟图片的对应关系就全乱了这是新手特别容易踩的坑。num_workers可以根据你的CPU核数适当调大数据读取的瓶颈在磁盘IO时能明显提速。3. 整体方案设计冻结模型还是微调3.1 三种主流的实现路线拿到DINOv2之后你的选择不是唯一的。根据你的数据量和任务难度大致有三条路。第一条路是“冻结backbone 训练线性分类头”也就是线性探针linear probing。做法是把DINOv2的所有参数冻结住对每张图片提取特征向量然后只训练一个逻辑回归层或单层全连接。这条路速度最快、最不容易过拟合在数据量只有几百张的情况下也能有不错的效果。它是我在这个项目里的首选。第二条路是“冻结backbone 训练MLP分类头”。跟第一条的区别是分类头不再是简单的线性层而是一个两层或三层的全连接网络中间加ReLU激活和Dropout。这能多学一些特征组合的模式适合线性分类头效果差一点点、但你又不希望动backbone的场景。第三条路是“全量微调”。把DINOv2的全部参数或者后面几层解冻用你的数据端到端地训练。这条路上限最高但是对你的数据量要求也最高而且训练时间会长很多。如果只有千把张图全量微调很容易过拟合。我的建议是先走第一条路把baseline打出来如果准确率不够再考虑微调后面几层。3.2 为什么线性探针往往就够了DINOv2学出来的特征有一个很强的特性线性可分性。也就是说不同类别的特征在向量空间里已经天然分得比较开了只是需要一层线性变换把坐标轴旋转到合适的角度。线性探针本质上就是在学这个变换。我在多个自定义数据集上测过只要你的类别不是特别抽象比如“开心”和“难过”这种情绪10000张训练图下线性探针的准确率往往能到95%以上。即便只有几百张图效果也远好于传统方法。这背后的原因就是DINOv2在无标注数据上见过海量的视觉模式它的特征空间已经非常平滑和规整线性分类头加在这上面就是顺水推舟。3.3 特征维度与模型参数速查模型参数量特征维度适用场景ViT-S/142100万384快速验证、资源受限ViT-B/148600万768通用推荐ViT-L/143亿1024高精度需求、显存充足ViT-g/1411亿1536极致效果、大规模任务特征维度直接影响后面线性层的输入大小和训练速度。768维的特征去训一个二分类逻辑回归计算量完全可以忽略不计。这个维度对于后续做图像检索、聚类等任务也都很合适。4. 核心实操加载DINOv2模型与提取特征4.1 用torch.hub加载预训练模型DINOv2的官方加载方式在facebookresearch/dinov2仓库里写得非常清楚。核心代码就一行import torch device torch.device(cuda if torch.cuda.is_available() else cpu) model torch.hub.load(facebookresearch/dinov2, dinov2_vitb14) model.eval() model.to(device)这里有个细节必须强调加载后一定记得调用model.eval()。因为模型内部有Dropout和LayerNorm等模块训练和推理模式下的行为不一样。不切换成eval模式提取的特征会有随机性导致同一个输入两次特征不一致后面训练分类器就会莫名其妙地不稳定。我第一版代码就吃了这个亏还以为是模型随机初始化的问题排查了半天才发现是忘了切换模式。torch.hub第一次加载时会去下载权重和代码如果网络环境不好下载会卡住。遇到这种情况可以手动把权重文件下载下来放到~/.cache/torch/hub/checkpoints/目录下文件名要跟它期望的一致。ViT-B/14的权重文件名是dinov2_vitb14_pretrain.pth。放到对应位置后torch.hub加载时会直接读缓存不会再走网络。4.2 特征提取的完整代码提取特征最推荐的做法是把图片整体过一遍模型取CLS token的输出。代码如下import torch import numpy as np from tqdm import tqdm def extract_features(model, dataloader, device): model.eval() features [] labels [] with torch.no_grad(): for images, targets in tqdm(dataloader): images images.to(device) outputs model(images) # shape: (batch_size, dim) features.append(outputs.cpu().numpy()) labels.append(targets.cpu().numpy()) features np.concatenate(features, axis0) labels np.concatenate(labels, axis0) return features, labels train_features, train_labels extract_features(model, train_loader, device) val_features, val_labels extract_features(model, val_loader, device) np.save(train_features.npy, train_features) np.save(train_labels.npy, train_labels) np.save(val_features.npy, val_features) np.save(val_labels.npy, val_labels)关于特征向量怎么取DINOv2的forward默认返回CLS token对应的向量。有些文章会建议把最后一层所有patch token做均值池化再跟CLS token拼接。我实测过在分类任务上直接用CLS token已经非常好了拼接其他token收益很小但特征维度变大了后续分类器训练和存储成本都会上升。所以常规场景下直接用CLS token就好。提取出来的特征建议先保存成npy文件。这样的好处是你后续尝试不同的分类头方案时不需要再重新过一遍模型直接在特征上做实验省时间也省显存。我通常会把这个特征提取脚本单独保存后续做特征可视化、聚类、检索都会再次用到它。4.3 先看特征有没有区分度训练分类器之前强烈建议先对特征做一个可视化看看不同类别的点在特征空间里是否真的能分开。这能帮你提前发现数据本身的问题。用t-SNE把768维的特征降到2维然后按标签染色from sklearn.manifold import TSNE import matplotlib.pyplot as plt tsne TSNE(n_components2, random_state42, perplexity30) features_2d tsne.fit_transform(train_features) plt.figure(figsize(8, 8)) colors [#e74c3c, #3498db] for label in np.unique(train_labels): idx train_labels label plt.scatter(features_2d[idx, 0], features_2d[idx, 1], s5, ccolors[label]) plt.savefig(tsne_features.png, dpi150)如果t-SNE图上两个类别明显聚成两团恭喜你后面分类器的效果基本不会差。如果两个类别完全混在一起那问题多半出在数据上比如标注错误、图片质量太差、类别本身难以区分等等。这时候不要急着调模型先回去看数据才是最明智的。我经常跟朋友说特征可视化是你和你的数据之间的照妖镜它比任何指标都能更直观地告诉你数据的真实情况。5. 核心实操训练分类头并评估效果5.1 线性探针最简单的分类器特征拿到手之后训练分类器就变成一个非常经典的小问题。最简单、最不容易翻车的方案是逻辑回归LogisticRegression。这里用scikit-learn的实现多分类和软标签都支持得很好from sklearn.linear_model import LogisticRegression from sklearn.metrics import accuracy_score, classification_report clf LogisticRegression(max_iter1000, C1.0) clf.fit(train_features, train_labels) val_pred clf.predict(val_features) val_acc accuracy_score(val_labels, val_pred) print(fValidation accuracy: {val_acc:.4f}) print(classification_report(val_labels, val_pred, target_namesval_dataset.classes))重点说一下LogisticRegression的几个关键参数。max_iter设大一点是因为特征维度768、样本量可能上万默认的100次迭代有时候不够收敛。C是正则化强度的倒数C越小正则化越强。如果训练集不大可以适当把C调小一点防止过拟合比如C0.1。你也可以用GridSearchCV快速搜索一个合适的C值不过一般0.1到10这个范围内效果差异不是特别大。训练时间上10000个样本、768维特征逻辑回归几乎是一眨眼就训完。这也是DINOv2方案让我觉得特别爽的地方前面重活累活模型都干完了后面的分类任务轻得跟羽毛一样。5.2 如果需要更高精度加一层MLP如果线性分类头在验证集上的表现差强人意你可以升级成一层带隐藏层的MLP。思路是把特征先映射到一个中间维度经过ReLU激活后再映射到类别数。我用PyTorch写了一个简单的版本import torch.nn as nn class MLPClassifier(nn.Module): def __init__(self, in_dim, hidden_dim512, num_classes2): super().__init__() self.net nn.Sequential( nn.Linear(in_dim, hidden_dim), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(hidden_dim, num_classes) ) def forward(self, x): return self.net(x) model_cls MLPClassifier(in_dim768, num_classes2).to(device) optimizer torch.optim.AdamW(model_cls.parameters(), lr1e-3) criterion nn.CrossEntropyLoss()训练时我自己习惯把特征张量转成PyTorch的Dataset然后用DataLoader按batch迭代。隐藏层维度取512在768维输入下是一个比较稳妥的选择。Dropout取0.3到0.5之间能有效缓解过拟合。MLP训练几十个epoch就够了千万别训太久因为特征本身已经很强了分类头参数又少过拟合很快会来敲门的。有朋友问我为什么不用SVM其实也可以用SVM在小样本下表现也挺好。但SVM的预测阶段不如逻辑回归那样方便输出概率在需要不确定性估计的场景下会比较别扭。逻辑回归在大多数情况下已经够用所以我默认先推它。5.3 评估指标不能只看准确率准确率是最直观的指标但如果你的类别不平衡准确率会骗人。比如95%的样本都是猫那模型把狗全部识别成猫也能有95%的准确率看起来挺好的实际上一无是处。所以我强烈建议在二分类场景下同时看混淆矩阵Confusion Matrix和每个类别的精确率、召回率和F1值。上面的classification_report已经把这些指标打出来了。混淆矩阵可以用下面的代码画出来from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay cm confusion_matrix(val_labels, val_pred) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labelsval_dataset.classes) disp.plot(cmapBlues) plt.savefig(confusion_matrix.png, dpi150)查看混淆矩阵能让你一眼发现模型到底在哪些类别上容易混淆。比如猫狗分类假如模型经常把白狗识别成猫可能是因为数据集中白猫太多或者白狗的样例太少。这时候针对性地补充对应类别的训练样本往往比调模型参数有用得多。5.4 我这次实验的实际效果用大约8000张训练图片、2000张验证图片做猫狗二分类ViT-B/14 逻辑回归的验证准确率是99.2%。这个结果几乎不需要调参第一次跑完就是这水平。后来我换了一个工业场景的数据集每类只有三四百张图DINOv2特征线性探针也能做到89%的准确率而传统ResNet迁移学习在那个数据集上大概只有82%左右。差距非常明显而且DINOv2方案的训练时间几乎可以忽略不计让我省出了大把时间去做数据分析和排错。6. 常见问题与排查技巧实录6.1 特征一模一样检查是不是忘了eval模式我见过好几个人跑来问说DINOv2提取的特征怎么会有随机性同一个输入两次特征都不一样。排查到最后基本都是同一句话模型没有调用model.eval()。PyTorch默认是训练模式Dropout层会随机屏蔽节点LayerNorm也会对batch统计量做修正这些都会导致输出不稳定。切记推理前调用model.eval()并且用torch.no_grad()包住前向传播。6.2 下载模型权重卡住了怎么办torch.hub加载DINOv2时首次需要下载数百MB的权重文件。网络状况不好时下载过程可能长时间没有进度。解决办法是手动下载权重放到缓存目录。以Linux为例缓存目录一般是~/.cache/torch/hub/checkpoints/把权重文件放到这里就行。Windows下是C:\Users\你的用户名.cache\torch\hub\checkpoints\。这样torch.hub加载时会直接命中缓存不再走网络。6.3 类别的图片数量严重不均衡如果某个类别的图片特别少线性分类头容易偏向样本量大的类别。我的处理思路有几个一是对少数类做简单的数据增强比如随机翻转、旋转、颜色抖动然后再提取特征二是把逻辑回归的class_weight设为balanced让损失函数自动给少数类加权三是实在不行就换用余弦相似度分类用每个类别的原型向量做最近邻判断这种思路在极度不均衡的少样本场景下特别管用。6.4 验证集上效果很好但新数据上一塌糊涂这是典型的过拟合或者数据分布不一致问题。常见原因有三个一是验证集和训练集来自同一批采集环境导致模型学到的是环境特征而不是物体本身特征所以换到新环境就失灵。二是数据增强不合理比如把训练图做了过度颜色扰动让模型学到了很奇怪的判断依据。三是你在做数据预处理和推理预处理时用了不一致的流程比如训练时缩放到224预测时直接resize成别的尺寸。建议严格统一预处理流程最好把预处理逻辑封装成一个函数训练和预测都用同一个。6.5 显存不够用显存不足通常不是模型本身的问题而是batch size设置太大了。ViT-B/14在224x224输入下单张图片的显存占用大约是几百MB。如果一张卡只有8GB显存batch size设32是稳妥的再大就可能爆显存。还嫌不够的话可以尝试半精度推理把模型转成float16model model.half()半精度能显著减少显存占用对特征提取的效果影响很小。需要注意的是输入图片也要转成半精度否则类型不匹配会直接报错。7. 从分类到更多玩法这套方案的扩展思路DINOv2特征的价值远不止训练一个分类头。我在实际项目中逐渐发现一旦你把图片变成特征向量很多以前觉得很麻烦的任务都变得轻巧了。比如图像检索。把整个图片库的特征都提取出来存成矩阵来了一个新查询图片提取特征后跟库里的向量做余弦相似度排序Top-K就是检索结果。这个方案不需要任何训练效果在大多数场景下已经能媲美专门的检索模型。再比如聚类。如果你手里有一堆完全没标注的图片想先把它们粗分成几组DINOv2特征搭配KMeans聚类就能做得不错。我在一个商品分类项目里就是用这种方式先做了预聚类然后再人工给每个簇打标签大大减少了标注工作量。还有少样本分类。每类只用五张图甚至一张图用特征向量取类内平均作为原型然后做最近邻分类也就是Prototypical Network的思路。DINOv2的特征质量决定了这个方案的上限实测比在ResNet特征上做同样的操作要高出一大截。回到这次的猫狗分类项目本身我的体会是DINOv2把“做图像分类”这件事的门槛大大降低了。以前你需要花大量时间在特征工程、网络结构设计和训练策略调优上现在你只需要把图片整理好过一遍预训练模型然后训一个简单的分类器就能得到不错的结果。这种“预训练特征 轻量分类头”的开发模式特别适合快速验证想法、做技术Demo、或者在小规模数据上落地业务。最后再分享一个小技巧特征和标签尽量保留成npy文件后续调分类头、做可视化、换评估指标都不需要重新提取特征非常节省时间。希望这篇记录能让你少走几步弯路如果遇到什么问题欢迎在评论区交流。
返回列表