? [2023-10-26] Vision Transformers Need Registers に続き、レジスタを備えた DINOv2 バックボーンを追加しました。
メタAIリサーチ、FAIR
マキシム・オクアブ、ティモシー・ダルセ、テオ・ムタカンニ、ユイ・V・ヴォー、マルク・ザフラニエツ、ワシル・ハリドフ、パトリック・ラバトゥ、アルマン・ジュラン、ピョートル・ボジャノフスキー
[ Paper #1
] Paper #2
] [ Blog
] [ Demo
] [ BibTeX
]
DINOv2 の PyTorch 実装と事前トレーニング済みモデル。詳細については、論文「 DINOv2: Learning Robust Visual features without Supervision」および「Vision Transformers Need Registers」を参照してください。
DINOv2 モデルは、さまざまなコンピューター ビジョン タスクの線形レイヤーと同じくらい単純な分類器で直接使用できる高性能の視覚機能を生成します。これらの視覚機能は堅牢であり、微調整を必要とせずに、複数のドメインにわたって良好に機能します。モデルは、ラベルや注釈を使用せずに、1 億 4,200 万枚の画像のデータセットで事前トレーニングされました。
モデル | # の パラメータ | と レジスター | イメージネット k-NN | イメージネット リニア | ダウンロード |
---|---|---|---|---|---|
ViT-S/14蒸留 | 21分 | 79.0% | 81.1% | バックボーンのみ | |
ViT-S/14蒸留 | 21分 | ✅ | 79.1% | 80.9% | バックボーンのみ |
ViT-B/14蒸留 | 86M | 82.1% | 84.5% | バックボーンのみ | |
ViT-B/14蒸留 | 86M | ✅ | 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 (モデルをロードするために必要な唯一の依存関係) をインストールしてください。 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 | マスク2元 | 線形、マルチスケール | 線形、マルチスケール |
トレーニングおよび評価コードには、PyTorch 2.0 および xFormers 0.0.18 に加えて、他の多くのサードパーティ パッケージが必要です。コードは指定されたバージョンでのみテストされており、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 (推奨) :
conda env create -f conda-extras.yaml
conda activate dinov2-extras
ピップ:
pip install -r requirements.txt -r requirements-extras.txt
データセットのルート ディレクトリには次の内容が含まれている必要があります。
<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 コード行を使用して (1 回だけ) 生成できます。
from dinov2 . data . datasets import ImageNet
for split in ImageNet . Split :
dataset = ImageNet ( split = split , root = "<ROOT>" , extra = "<EXTRA>" )
dataset . dump_extra ()
ルート ディレクトリと追加ディレクトリは別個のディレクトリである必要はないことに注意してください。
データセット クラスをローカル設定に合わせて調整してください。
dinov2
パッケージを Python モジュール検索パスに含める必要があります。つまり、実行するコマンドの先頭にPYTHONPATH=.
。
submitit を使用して、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
フォルダーに保存します。
submitit を使用して、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
フォルダーに保存します。
トレーニング コードは定期的に教師の重みを保存します。モデルを評価するには、単一ノードで次の評価を実行します。
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 >
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 >
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 >
コミュニティがモデルとコードを活用できるように、いくつかのノートブックが提供されています。
DINOv2 コードとモデルの重みは、Apache License 2.0 に基づいてリリースされます。詳細については、「ライセンス」を参照してください。
貢献と行動規範をご覧ください。
このリポジトリが役立つと思われる場合は、スターと引用を付けることを検討してください ?:
@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}
}