DINOv2 - 无监督学习鲁棒视觉特征
本文翻译整理自:https://github.com/facebookresearch/dinov2
文章目录
- 一、关于 DINOv2
- 相关链接资源
- 关键功能特性
- 二、预训练模型
- 预训练骨架网络
- 通过 PyTorch Hub 加载预训练模型
- 预训练分类头 - ImageNet
- 预训练头 - 深度估计
- 预训练头 - 语义分割
- 三、安装
- 1、推荐安装方式(conda)
- 2、pip 安装方式
- 3、密集任务额外依赖
- conda 方式(推荐)
- pip 方式
- 四、数据准备
- 1、ImageNet-1k 数据集
- 2、ImageNet-22k 数据集
- 五、训练
- 1、快速设置:在 ImageNet-1k 上训练 DINOv2 ViT-L/16
- 2、完整设置:在 ImageNet-22k 上训练 DINOv2 ViT-L/14
- 六、评估
- 1、ImageNet-1k 上的 k-NN 分类
- 2、ImageNet-1k 上的逻辑回归分类
- 3、带数据增强的 ImageNet-1k 线性分类
- 4、预训练模型评估
- 七、Notebooks
一、关于 DINOv2
DINOv2 是 Meta AI Research (FAIR) 团队开发的视觉特征提取模型,能够在不使用任何标签或标注的情况下,从 1.42 亿张图像中学习高性能的视觉特征。
这些特征可以直接用于各种计算机视觉任务,只需配合简单的线性分类器即可获得优异表现。
主要贡献者:Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Marc Szafraniec, Vasil Khalidov, Patrick Labatut, Armand Joulin, Piotr Bojanowski
相关链接资源
- github : https://github.com/facebookresearch/dinov2
- 官网:https://ai.facebook.com/research/
- 官方文档:https://github.com/facebookresearch/dinov2
- Paper :
- https://arxiv.org/abs/2304.07193
- https://arxiv.org/abs/2309.16588 注册机制
- Demo/在线试用:https://dinov2.metademolab.com
- Blog : https://ai.facebook.com/blog/dino-v2-computer-vision-self-supervised-learning/
- BibTeX : 见文末引用部分
- License : Apache License 2.0
关键功能特性
- 完全无监督预训练,无需任何标注数据
- 生成的视觉特征具有强大的跨领域泛化能力
- 支持多种视觉任务:分类、深度估计、语义分割等
- 提供不同规模的模型:ViT-S/14、ViT-B/14、ViT-L/14、ViT-g/14
- 2023-10-26 新增:支持带寄存器的 DINOv2 骨干网络
二、预训练模型
预训练骨架网络
模型 | 参数量 | 带寄存器 | ImageNet k-NN | ImageNet linear | 下载 |
---|---|---|---|---|---|
ViT-S/14 distilled | 21 M | ❌ | 79.0% | 81.1% | backbone only |
ViT-S/14 distilled | 21 M | ✅ | 79.1% | 80.9% | backbone only |
ViT-B/14 distilled | 86 M | ❌ | 82.1% | 84.5% | backbone only |
ViT-B/14 distilled | 86 M | ✅ | 82.0% | 84.6% | backbone only |
ViT-L/14 distilled | 300 M | ❌ | 83.5% | 86.3% | backbone only |
ViT-L/14 distilled | 300 M | ✅ | 83.8% | 86.7% | backbone only |
ViT-g/14 | 1,100 M | ❌ | 83.5% | 86.5% | backbone only |
ViT-g/14 | 1,100 M | ✅ | 83.7% | 87.1% | backbone only |
通过 PyTorch Hub 加载预训练模型
请先按照官方指南安装 PyTorch(加载模型所需的唯一依赖项),强烈建议安装支持 CUDA 的 PyTorch 版本。
import torch# DINOv2
dinov2_vits14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
dinov2_vitb14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
dinov2_vitl14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
dinov2_vitg14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14')# 带寄存器的 DINOv2
dinov2_vits14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14_reg')
dinov2_vitb14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
dinov2_vitl14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg')
dinov2_vitg14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_reg')
预训练分类头 - ImageNet
骨干网络 | 带寄存器 | ImageNet 分类头下载 |
---|---|---|
ViT-S/14 distilled | ❌ | linear head (1层, 4层) |
ViT-S/14 distilled | ✅ | linear head (1层, 4层) |
ViT-B/14 distilled | ❌ | linear head (1层, 4层) |
ViT-B/14 distilled | ✅ | linear head (1层, 4层) |
ViT-L/14 distilled | ❌ | linear head (1层, 4层) |
ViT-L/14 distilled | ✅ | linear head (1层, 4层) |
ViT-g/14 | ❌ | linear head (1层, 4层) |
ViT-g/14 | ✅ | linear head (1层, 4层) |
完整分类器模型可通过 PyTorch Hub 加载:
import torch# DINOv2
dinov2_vits14_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14_lc')
dinov2_vitb14_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_lc')
dinov2_vitl14_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_lc')
dinov2_vitg14_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_lc')# 带寄存器的 DINOv2
dinov2_vits14_reg_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14_reg_lc')
dinov2_vitb14_reg_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg_lc')
dinov2_vitl14_reg_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg_lc')
dinov2_vitg14_reg_lc = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_reg_lc')
预训练头 - 深度估计
骨干网络 | NYUd 下载头 | KITTI 下载头 |
---|---|---|
ViT-S/14 distilled | linear (1层, 4层), DPT | linear (1层, 4层), DPT |
ViT-B/14 distilled | linear (1层, 4层), DPT | linear (1层, 4层), DPT |
ViT-L/14 distilled | linear (1层, 4层), DPT | linear (1层, 4层), DPT |
ViT-g/14 | linear (1层, 4层), DPT | linear (1层, 4层), DPT |
预训练头 - 语义分割
骨干网络 | ADE20K 模型下载 | ADE20K 头下载 | VOC2012 头下载 |
---|---|---|---|
ViT-S/14 distilled | - | linear, multi-scale | linear, multi-scale |
ViT-B/14 distilled | - | linear, multi-scale | linear, multi-scale |
ViT-L/14 distilled | - | linear, multi-scale | linear, multi-scale |
ViT-g/14 | Mask2Former | linear, multi-scale | linear, multi-scale |
三、安装
训练和评估代码需要 PyTorch 2.0 和 xFormers 0.0.18 以及其他多个第三方包。
请注意,代码仅经过指定版本的测试,且需要 Linux 环境。
请按以下说明设置所有训练和评估所需的依赖项:
1、推荐安装方式(conda)
克隆仓库后,使用提供的环境定义文件创建并激活 dinov2
conda 环境:
conda env create -f conda.yaml
conda activate dinov2
2、pip 安装方式
克隆仓库后,使用提供的 requirements.txt
安装依赖项:
pip install -r requirements.txt
3、密集任务额外依赖
对于深度估计和语义分割等密集任务,需要额外安装特定版本的 mmcv
和 mmsegmentation
:
conda 方式(推荐)
conda env create -f conda-extras.yaml
conda activate dinov2-extras
pip 方式
pip install -r requirements.txt -r requirements-extras.txt
四、数据准备
1、ImageNet-1k 数据集
数据集根目录应包含以下内容:
<ROOT>/test/ILSVRC2012_test_00000001.JPEG
<ROOT>/test/[..]
<ROOT>/test/ILSVRC2012_test_00100000.JPEG
<ROOT>/train/n01440764/n01440764_10026.JPEG
<ROOT>/train/[...]
<ROOT>/train/n15075141/n15075141_9993.JPEG
<ROOT>/val/n01440764/ILSVRC2012_val_00000293.JPEG
<ROOT>/val/[...]
<ROOT>/val/n15075141/ILSVRC2012_val_00049174.JPEG
<ROOT>/labels.txt
额外目录需要包含以下元数据文件:
<EXTRA>/class-ids-TRAIN.npy
<EXTRA>/class-ids-VAL.npy
<EXTRA>/class-names-TRAIN.npy
<EXTRA>/class-names-VAL.npy
<EXTRA>/entries-TEST.npy
<EXTRA>/entries-TRAIN.npy
<EXTRA>/entries-VAL.npy
可通过以下 Python 代码生成这些元数据文件:
from dinov2.data.datasets import ImageNetfor split in ImageNet.Split:dataset = ImageNet(split=split, root="<ROOT>", extra="<EXTRA>")dataset.dump_extra()
注:根目录和额外目录可以是同一个目录。
2、ImageNet-22k 数据集
请根据本地设置调整 dataset class。
警告:执行后续训练和评估命令时,需将 dinov2
包包含在 Python 模块搜索路径中,即在命令前添加 PYTHONPATH=.
。
五、训练
1、快速设置:在 ImageNet-1k 上训练 DINOv2 ViT-L/16
在 SLURM 集群环境(4 个 A100-80GB 节点,32 个 GPU)上运行训练:
python dinov2/run/train/train.py \--nodes 4 \--config-file dinov2/configs/train/vitl16_short.yaml \--output-dir <PATH/TO/OUTPUT/DIR> \train.dataset_path=ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
训练时间约 1 天,最终 checkpoint 在 k-NN 评估中应达到 81.6%,线性评估中达到 82.9%。
训练代码每 12500 次迭代会将教师权重保存在 eval
文件夹中供评估使用。
2、完整设置:在 ImageNet-22k 上训练 DINOv2 ViT-L/14
在 SLURM 集群环境(12 个 A100-80GB 节点,96 个 GPU)上运行训练:
python dinov2/run/train/train.py \--nodes 12 \--config-file dinov2/configs/train/vitl14.yaml \--output-dir <PATH/TO/OUTPUT/DIR> \train.dataset_path=ImageNet22k:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
训练时间约 3.3 天,最终 checkpoint 在 k-NN 评估中应达到 82.0%,线性评估中达到 84.5%。
训练代码每 12500 次迭代会将教师权重保存在 eval
文件夹中供评估使用。
六、评估
训练代码会定期保存教师权重。要在单节点上评估模型,请运行以下命令:
1、ImageNet-1k 上的 k-NN 分类
python dinov2/run/eval/knn.py \--config-file <PATH/TO/OUTPUT/DIR>/config.yaml \--pretrained-weights <PATH/TO/OUTPUT/DIR>/eval/training_24999/teacher_checkpoint.pth \--output-dir <PATH/TO/OUTPUT/DIR>/eval/training_24999/knn \--train-dataset ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \--val-dataset ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
2、ImageNet-1k 上的逻辑回归分类
python dinov2/run/eval/log_regression.py \--config-file <PATH/TO/OUTPUT/DIR>/config.yaml \--pretrained-weights <PATH/TO/OUTPUT/DIR>/eval/training_24999/teacher_checkpoint.pth \--output-dir <PATH/TO/OUTPUT/DIR>/eval/training_24999/logreg \--train-dataset ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \--val-dataset ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
3、带数据增强的 ImageNet-1k 线性分类
python dinov2/run/eval/linear.py \--config-file <PATH/TO/OUTPUT/DIR>/config.yaml \--pretrained-weights <PATH/TO/OUTPUT/DIR>/eval/training_24999/teacher_checkpoint.pth \--output-dir <PATH/TO/OUTPUT/DIR>/eval/training_24999/linear \--train-dataset ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \--val-dataset ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
4、预训练模型评估
我们发布了不同模型的评估权重:
model | with registers | ImageNet top-1 | linear evaluation |
---|---|---|---|
ViT-S/14 distilled | ❌ | 81.1% | linear head weights |
ViT-S/14 distilled | ✅ | 80.8% | linear head weights |
ViT-B/14 distilled | ❌ | 84.5% | linear head weights |
ViT-B/14 distilled | ✅ | 84.4% | linear head weights |
ViT-L/14 distilled | ❌ | 86.3% | linear head weights |
ViT-L/14 distilled | ✅ | 86.5% | linear head weights |
ViT-g/14 | ❌ | 86.5% | linear head weights |
ViT-g/14 | ✅ | 87.0% | linear head weights |
使用以下命令评估预训练模型在 ImageNet-1k 上的性能:
python dinov2/run/eval/linear.py \--config-file dinov2/configs/eval/vitg14_pretrain.yaml \--pretrained-weights https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_pretrain.pth \--train-dataset ImageNet:split=TRAIN:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET> \--val-dataset ImageNet:split=VAL:root=<PATH/TO/DATASET>:extra=<PATH/TO/DATASET>
七、Notebooks
我们提供了一些 Notebook 帮助社区使用模型和代码:
- 深度估计 - 如何通过 mmcv 加载和使用与骨干网络匹配的深度头
- 语义分割 - 如何通过 mmcv 加载和使用与骨干网络匹配的分割头,以及如何加载和使用在 ADE20K 上训练的基于 Mask2Former 的分割模型
伊织 xAI 2025-04-25(五)