DINOv2: 監督なしでロバストな視覚的特徴を学習する【セグメンテーション】

メタAI研究、FAIR

Maxime Oquab、Timothée Darcet、Théo Moutakanni、Huy Vo、Marc Szafraniec、Vasil Khalidov、Patrick Labatut、Armand Joulin、Piotr Bojanowski

[ Paper] [ Blog] [ Demo] [ BibTeX]

DINOv2 の PyTorch 実装と事前トレーニング済みモデル。詳細については、次の論文を参照してください: DINOv2: Learning Robust Visual Features without Supervision .

DINOv2 モデルは、さまざまなコンピューター ビジョン タスクの線形レイヤーと同じくらい単純な分類器で直接使用できる高性能の視覚的特徴を生成します。これらの視覚的機能は堅牢で、微調整を必要とせずにドメイン全体で適切に機能します。モデルは、ラベルや注釈を使用せずに、1 億 4,200 万枚の画像のデータセットで事前トレーニングされました。

 ビデオ参照 + dinov2.mp4 
RGB 値にマッピングされた、すべてのフレームのパッチ機能の最初の 3 つの主成分の視覚化。
事前トレーニング済みモデル
モデル	
パラメータの数	ImageNet
k-NN	ImageNet
線形	ダウンロード
ViT-S/14蒸留	21M	79.0%	81.1%	バックボーンのみ
ViT-B/14 蒸留	86M	82.1%	84.5%	バックボーンのみ
ViT-L/14蒸留	300M	83.5%	86.3%	バックボーンのみ
ViT-g/14	1,100M	83.5%	86.5%	バックボーンのみ
PyTorch Hub による事前トレーニング済みモデル
こちらの手順に従って、PyTorch と torchvision の依存関係をインストールしてください (必要な依存関係はこれらだけです)。CUDA をサポートする PyTorch と torchvision の両方をインストールすることを強くお勧めします。

対応するモデル カードは [ ] ファイルにありますMODEL_CARD.md。

import torch

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')
インストール
トレーニングと評価のコードには、PyTorch 2.0 と xFormers 0.0.18、およびその他の多くのサードパーティ パッケージが必要です。トレーニングと評価に必要なすべての依存関係をセットアップするには、以下の手順に従ってください。

conda (推奨) -dinov2提供された環境定義を使用して、conda 環境を作成およびアクティブ化します。

conda env create -f conda.yaml
conda activate dinov2
pip - 提供されたものを使用しrequirements.txtて依存関係をインストールします。

pip install -r requirements.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
ImageNet-22k の場合は、それに応じて Dataset オブジェクトを調整してください。

トレーニング
迅速なセットアップ: ImageNet-1k での DINOv2 ViT-L/16 のトレーニング
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フォルダーに保存します。

長いセットアップ: ImageNet-22k での DINOv2 ViT-L/14 のトレーニング
submitit を使用して、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 日で、結果のチェックポイントは 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>
さまざまなモデルの評価から重みを解放します。

モデル	ImageNet
トップ 1	線形評価
ViT-S/14蒸留	81.1%	リニアヘッドウェイト
ViT-B/14 蒸留	84.5%	リニアヘッドウェイト
ViT-L/14蒸留	86.3%	リニアヘッドウェイト
ViT-g/14	86.5%	リニアヘッドウェイト
提供された事前トレーニング済みモデルの重みのパフォーマンスは、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>
ライセンス
このリポジトリとモデルは、 LICENSEファイルにあるように、CC-BY-NC の下でリリースされます。

貢献する
貢献と行動規範を参照してください。