首页>编程相关>其他源码
下载

？ [2023-10-26]根据 Vision Transformers Need Registers 添加了带寄存器的 DINOv2 主干。

DINOv2：在没有监督的情况下学习鲁棒的视觉特征

元人工智能研究，FAIR

马克西姆·奥夸布、蒂莫西·达尔塞、西奥·穆塔卡尼、Huy V. Vo、马克·萨弗拉涅克、瓦西尔·哈利多夫、帕特里克·拉巴图、阿曼德·朱林、皮奥特·博亚诺夫斯基

[ Paper #1 ] Paper #2 ] [ Blog ] [ Demo ] [ BibTeX ]

DINOv2 的 PyTorch 实现和预训练模型。详情参见论文： DINOv2: Learning Robust Visual Features without Supervision和Vision Transformers Need Registers 。

DINOv2 模型产生高性能的视觉特征，可以直接与各种计算机视觉任务中的线性层一样简单的分类器使用；这些视觉功能非常强大，并且跨领域表现良好，无需任何微调。这些模型在 1.42 亿张图像的数据集上进行预训练，未使用任何标签或注释。

视频参考+dinov2.mp4
所有帧的补丁特征的前三个主要成分的可视化，映射到 RGB 值。

预训练模型

模型＃ 的
参数		和
寄存器		图像网
k-NN		图像网
线性		下载
ViT-S/14 蒸馏21米79.0% 81.1%仅骨干
ViT-S/14 蒸馏21米✅ 79.1% 80.9%仅骨干
ViT-B/14 蒸馏86米82.1% 84.5%仅骨干
ViT-B/14 蒸馏86米✅ 82.0% 84.6%仅骨干
ViT-L/14 蒸馏300M 83.5% 86.3%仅骨干
ViT-L/14 蒸馏300M ✅ 83.8% 86.7%仅骨干
ViT-g/14 1,100M 83.5% 86.5%仅骨干
ViT-g/14 1,100M ✅ 83.7% 87.1%仅骨干

预训练主干（通过 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 with registers
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' )

预训练头 - 图像分类

骨干和
寄存器		下载
图像网
ViT-S/14 蒸馏直线头（1层、4层）
ViT-S/14 蒸馏✅直线头（1层、4层）
ViT-B/14 蒸馏直线头（1层、4层）
ViT-B/14 蒸馏✅直线头（1层、4层）
ViT-L/14 蒸馏直线头（1层、4层）
ViT-L/14 蒸馏✅直线头（1层、4层）
ViT-g/14直线头（1层、4层）
ViT-g/14 ✅直线头（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 with registers
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' )

预训练头 - 深度估计

骨干下载头
纽约大学基蒂
ViT-S/14 蒸馏线性（1层、4层）、DPT线性（1层、4层）、DPT
ViT-B/14 蒸馏线性（1层、4层）、DPT线性（1层、4层）、DPT
ViT-L/14 蒸馏线性（1层、4层）、DPT线性（1层、4层）、DPT
ViT-g/14线性（1层、4层）、DPT线性（1层、4层）、DPT

预训练头 - 语义分割

骨干下载模型下载头
ADE20K ADE20K VOC2012
ViT-S/14 蒸馏线性、多尺度线性、多尺度
ViT-B/14 蒸馏线性、多尺度线性、多尺度
ViT-L/14 蒸馏线性、多尺度线性、多尺度
ViT-g/14 Mask2Former线性、多尺度线性、多尺度

安装

培训和评估代码需要 PyTorch 2.0 和 xFormers 0.0.18 以及许多其他第 3 方软件包。请注意，代码仅在指定版本上进行了测试，并且还需要 Linux 环境。要设置培训和评估所需的所有依赖项，请按照以下说明操作：

conda （推荐） - 克隆存储库，然后使用提供的环境定义创建并激活dinov2 conda 环境： 

conda env create -f conda.yaml
conda activate dinov2

pip - 克隆存储库，然后使用提供的requirements.txt安装依赖项： 

pip install -r requirements.txt

对于密集任务（深度估计和语义分割），还有额外的依赖项（ mmcv和mmsegmentation的特定版本），这些依赖项在extras依赖项规范中捕获：

康达（推荐） ： 

conda env create -f conda-extras.yaml
conda activate dinov2-extras

点： 

pip install -r requirements.txt -r requirements-extras.txt

数据准备

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 ImageNet

for split in ImageNet . Split :
    dataset = ImageNet ( split = split , root = "<ROOT>" , extra = "<EXTRA>" )
    dataset . dump_extra ()

请注意，根目录和额外目录不必是不同的目录。

ImageNet-22k

请调整数据集类以匹配您的本地设置。


️要执行下一节中提供的命令进行训练和评估， dinov2包应包含在 Python 模块搜索路径中，即只需在要运行的命令前加上PYTHONPATH=. 。

训练

快速设置：在 ImageNet-1k 上训练 DINOv2 ViT-L/16

在 SLURM 集群环境中的 4 个 A100-80GB 节点（32 个 GPU）上运行 DINOv2 训练并提交： 

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 天，生成的检查点在 k-NN 评估上应达到 81.6%，在线性评估上应达到 82.9%。

训练代码每12500次迭代都会将教师的权重保存在eval文件夹中以进行评估。

长设置：在 ImageNet-22k 上训练 DINOv2 ViT-L/14

在 SLURM 集群环境中的 12 个 A100-80GB 节点（96 个 GPU）上运行 DINOv2 训练并提交： 

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 天，所得检查点在 k-NN 评估上应达到 82.0%，在线性评估上应达到 84.5%。

训练代码每12500次迭代都会将教师的权重保存在eval文件夹中以进行评估。

评估

训练代码定期保存教师权重。为了评估模型，请在单个节点上运行以下评估：

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 >

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 >

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 >

我们通过评估不同模型来释放权重：

模型和
寄存器		图像网
前1名		线性评估
ViT-S/14 蒸馏81.1%线性头配重
ViT-S/14 蒸馏✅ 80.8%线性头配重
ViT-B/14 蒸馏84.5%线性头配重
ViT-B/14 蒸馏✅ 84.4%线性头配重
ViT-L/14 蒸馏86.3%线性头配重
ViT-L/14 蒸馏✅ 86.5%线性头配重
ViT-g/14 86.5%线性头配重
ViT-g/14 ✅ 87.0%线性头配重

所提供的预训练模型权重的性能可以在 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 >

笔记本电脑

提供了一些笔记本来帮助社区利用模型和代码：

  • 深度估计 - 如何通过 mmcv 加载和使用深度头以及匹配的骨干网
  • 语义分割 - 如何通过 mmcv 加载和使用与匹配主干相结合的分割头，以及如何加载和使用在 ADE20K 上训练的基于 Mask2Former 的分割模型

执照

DINOv2 代码和模型权重在 Apache License 2.0 下发布。有关更多详细信息，请参阅许可证。

贡献

请参阅贡献和行为准则。

引用 DINOv2

如果您发现此存储库有用，请考虑给予星号和引用？： 

 @misc{oquab2023dinov2,
  title={DINOv2: Learning Robust Visual Features without Supervision},
  author={Oquab, Maxime and Darcet, Timothée and Moutakanni, Theo and Vo, Huy V. and Szafraniec, Marc and Khalidov, Vasil and Fernandez, Pierre and Haziza, Daniel and Massa, Francisco and El-Nouby, Alaaeldin and Howes, Russell and Huang, Po-Yao and Xu, Hu and Sharma, Vasu and Li, Shang-Wen and Galuba, Wojciech and Rabbat, Mike and Assran, Mido and Ballas, Nicolas and Synnaeve, Gabriel and Misra, Ishan and Jegou, Herve and Mairal, Julien and Labatut, Patrick and Joulin, Armand and Bojanowski, Piotr},
  journal={arXiv:2304.07193},
  year={2023}
}

 @misc{darcet2023vitneedreg,
  title={Vision Transformers Need Registers},
  author={Darcet, Timothée and Oquab, Maxime and Mairal, Julien and Bojanowski, Piotr},
  journal={arXiv:2309.16588},
  year={2023}
}
展开
附加信息
相关应用
为您推荐
相关资讯 全部