首頁>編程相關>其他源碼
下載

？ [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}
}
展開
附加信息
相關應用
爲您推薦
相關資訊 全部