最近的研究表明，在語義分割和目標檢測等計算機視覺任務中，大型 transformer 模型可以實現或提高 SOTA 。然而，與只能使用標準公共數據集的卷積網絡模型不同，它需要一個更大的專有數據集。
VOLO 模型體系結構
新加坡 SEA AI 實驗室最近的項目 VOLO （ Vision Outlooker ）展示了一種高效且可擴展的 Vision transformer 模式體系結構，該體系結構僅使用 ImageNet-1K 數據集就大大縮小了差距。
VOLO 引入了一種新穎的 outlook attention ，并提出了一種簡單而通用的架構，稱為 Vision Outlooker 。與自我關注不同，自我關注側重于粗略級別的全局依賴關系建模， outlook 關注有效地將更精細級別的功能和上下文編碼為標記。這對識別性能極為有利，但在很大程度上被自我注意所忽視。
實驗表明， VOLO 在 ImageNet-1K 分類上達到了 87.1% 的 top-1 精度，這是第一個在這個競爭基準上超過 87% 精度的模型，無需使用任何額外的訓練數據。
此外，經過預訓練的 VOLO 可以很好地轉移到下游任務，例如語義切分。
| Settings | LV-ViT | CaiT | NFNet-F6 | NFNNet-F5 | VOLO-D5 |
| Test Resolution | 448×448 | 448×448 | 576×576 | 544×544 | 448×448/512×512 |
| Model Size | 140M | 356M | 438M | 377M | 296M |
| Computations | 157B | 330B | 377B | 290B | 304B/412B |
| Architecture | Vision Transformer | Vision Transformer | Convolutions | Convolutions | VOLO |
| Extra Augmentations | Token Labeling | Knowledge Distill | SAM | SAM+augmult | Token Labeling |
| ImageNet Top-1 Acc. | 86.4 | 86.5 | 86.5 | 86.8 | 87.0/87.1 |
盡管 VOLO 模型顯示出了出色的計算效率，但訓練 SOTA 性能模型并非易事。
在這篇文章中，我們將介紹我們在 NVIDIA DGX SuperPOD 上基于 NVIDIA ML 軟件堆棧和 Infiniband 群集技術培訓 VOLO 模型所獲得的技術和經驗。
培訓方法
培訓 VOLO 模型需要考慮培訓策略、基礎設施和配置規劃。在本節中，我們將討論此解決方案中應用的一些技術。
培訓策略
始終使用原始 ImageNet 樣本質量數據訓練模型，并在細粒度上執行神經網絡（ NN ）架構搜索，使理論上的研究更加鞏固。然而，這需要計算資源預算的很大一部分。
在這個項目的范圍內，我們采用了一種粗粒度的訓練方法，它不像細粒度的方法那樣能夠訪問盡可能多的神經網絡體系結構。然而，它能夠以更少的時間和更低的資源預算顯示 EIOF 。在這種替代策略中，我們首先使用分辨率較低的圖像樣本訓練潛在的神經網絡候選，然后使用高分辨率圖像進行微調。
在早期的工作中，這種方法在降低邊際模型性能損失的計算成本方面被證明是有效的。
基礎設施
實際上，我們在本次培訓中使用了兩種類型的集群：
- 一個用于基礎模型預訓練，它是一個基于 NVIDIA DGX A100 的 DGX 吊艙，由使用 NVIDIA Mellanox HDR Infiniband 網絡集群的 5 個 NVIDIA DGX A100 系統組成。
- 一個用于微調，即 NVIDIA DGX SuperPOD ，由 DGX A100 系統和 NVIDIA Mellanox HDR Infiniband 網絡組成。
軟件基礎設施在這一過程中也發揮了重要作用。圖 2 顯示，除了基礎的標準深度學習優化 CUDA 庫（如 cuDNN 和 cuBLAS ）外，我們還廣泛利用 NCCL 、 enroot 、 PyXis 、 APEX 和 DALI 來實現培訓性能的亞線性可擴展性。
DGX A100 POD 集群主要用于使用較小尺寸圖像樣本的基礎模型預訓練。這是因為基本模型預訓練的內存限制較少，可以利用 NVIDIA A100 GPU 的計算能力優勢。
相比之下，微調是在 NVIDIA DGX-2 的 NVIDIA DGX SuperPOD 上執行的，因為微調過程使用更大的圖像，每臺計算能力需要更多的內存。
培訓配置
需要引入句子
| ? | D1 | D2 | D3 | D4 | D5 |
| MLP Ratio | 3 | 3 | 3 | 3 | 4 |
| Optimizer | AdamW | ||||
| LR Scaling | LR = LRbase x Batch_Size/1024, where LRbase=8.0e-4 | ||||
| Weight Decay | 5e-2 | ||||
| LRbase | 1.6e-2 | 1e-3 | 1e-3 | 1e-3 | 1e-4 |
| Stochastic Depth Rate | 0.1 | 0.2 | 0.5 | 0.5 | 0.75 |
| Crop Ratio | 0.96 | 0.96 | 0.96 | 1.15 | 1.15 |
表 2 ：。模型設置（對于所有模型，批大小設置為 1024 ）
我們在 ImageNet 數據集上評估了我們提出的 VOLO 模型。在培訓期間，沒有使用額外的培訓數據。我們的代碼基于 PyTorch 、令牌標記工具箱和 PyTorch 圖像模型（ timm ）。我們使用帶有標記的 LV-ViT-S 模型作為基線。
安裝說明
- 我們使用了 AdamW 優化器和線性學習率縮放策略 LR = LR基礎x Batch \ u 大小/ 1024 和 5 × 10 ? 2 先前工作建議的重量衰減率，表 3 中給出了所有 VOLO 模型的 LRbase 。
- 使用隨機深度。
- 我們在 ImageNet 數據集上訓練了 300 個時代的模型。
- 對于數據擴充方法，我們使用 CutOut 、 RandAug 和 MixToken 的標記目標。
- 我們沒有使用 MixUp 或 CutMix ，因為它們與 MixToken 沖突。
訓練前
在本節中，我們以 VOLO-D5 為例來演示如何訓練模型。
圖 3 顯示，使用單個 DGX A100 的 VOLO-D5 的訓練吞吐量約為 500 圖像/秒。據估計，完成一個完整的預訓練周期大約需要 170 個小時，這需要使用 ImageNet-1K 進行 300 個階段。這相當于 100 萬張圖片的一周。
為了加快速度，基于一個由五個 DGX A100 節點組成的簡單參數服務器架構集群，我們大致實現了 2100 個圖像/秒的吞吐量，這可以將預訓練時間減少到約 52 小時。
VOLO-D5 模型預訓練可以使用以下代碼示例在單個節點上啟動：
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet \ --model volo_d5 --img-size 224 \ -b 44 --lr 1.0e-4 --drop-path 0.75 --apex-amp \ --token-label --token-label-size 14 --token-label-data /path/to/token_label_data
對于 MNMG 培訓案例，它需要將培訓集群詳細信息作為命令行輸入的一部分。首先，我們根據節點和集群架構設置 CPU 、 MEM 、 IB 綁定。預訓練階段的集群是 DGX A100 POD ，每個 CPU 插槽有四個 NUMA 域，每個 A100 GPU 有一個 IB 端口，因此我們將每個列組綁定到 NUMA 節點中距離其 GPU 最近的所有 CPU 核。
- 對于內存綁定，我們將每個列組綁定到最近的 NUMA 節點。
- 對于 IB 綁定，我們為每個 GPU 綁定一個 IB 卡，或者盡可能接近這樣的設置。
由于 VOLO 模型培訓基于 PyTorch ，并且簡單地利用了默認的 PyTorch 分布式培訓方法，因此我們的多節點多 GPU 培訓基于一個簡單的參數服務器架構，該架構適合 NVIDIA DGX SuperPOD 的 fat 樹網絡拓撲。
為了簡化調度，分配節點列表中的第一個節點始終用作參數服務器和工作節點，而所有其他節點都是工作節點。為了避免潛在的存儲 I / O 開銷，數據集、所有代碼、中間/里程碑檢查點和結果都保存在一個基于 DDN 的高性能分布式存儲后端。它們通過 100G NVIDIA Mellanox EDR Infiniband 網絡裝載到所有工作節點。
為了加速數據預處理和流水線數據加載， NVIDIA DALI 配置為每個 GPU 進程使用一個專用數據加載程序。
微調
使用以下代碼示例，在單個節點上運行 VOLO-D5 模型微調非常簡單：
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet \ --model volo_d5 --img-size 512 \ -b 4 --lr 2.3e-5 --drop-path 0.5 --apex-amp --epochs 30 \ --weight-decay 1.0e-8 --warmup-epochs 5 --ground-truth \ --token-label --token-label-size 24 --token-label-data /path/to/token_label_data \ --finetune /path/to/pretrained_224_volo_d5/
如前所述，由于用于微調的圖像大小遠遠大于預訓練階段使用的圖像大小，因此必須相應地減小批量大小。將工作負載放入 GPU 內存中，這使得進一步擴展訓練到更大數量的 GPU 并行任務是必須的。
大多數微調配置類似于預訓練階段。
結論
在這篇文章中，我們展示了在大規模人工智能超級計算機上訓練 SOTA 大規模視覺 transformer 模型（如 VOLO \ u D5 ）的主要技術和程序，如基于 NVIDIA DGX A100 的 DGX SuperPOD 。經過訓練的 VOLO \ u D5 模型在圖像分類模型排名中取得了最佳的 Top-1 精度，無需使用 ImageNet-1k 數據集以外的任何其他數據。
這項工作的代碼資源（包括用于運行實驗的 Docker 映像和 Slurm 調度程序腳本）在 sail-sg/volo GitHub repo 中是開源的，以便將來可以在 VOLO \ u D5 上進行更廣泛的研究。有關更多信息，請參閱 VOLO ：視覺識別的視覺觀察家 。
未來，我們希望進一步擴展這項工作，以培訓更智能、自我監督、更大規模的模型，使用更大的公共數據集和更現代化的基礎設施，例如， NVIDIA DGX SuperPOD 和 NVIDIA H100 GPU。
?
