最近的研究表明，在語(yǔ)義分割和目標(biāo)檢測(cè)等計(jì)算機(jī)視覺(jué)任務(wù)中，大型 transformer 模型可以實(shí)現(xiàn)或提高 SOTA 。然而，與只能使用標(biāo)準(zhǔn)公共數(shù)據(jù)集的卷積網(wǎng)絡(luò)模型不同，它需要一個(gè)更大的專有數(shù)據(jù)集。

VOLO 模型體系結(jié)構(gòu)

新加坡 SEA AI 實(shí)驗(yàn)室最近的項(xiàng)目 VOLO （ Vision Outlooker ）展示了一種高效且可擴(kuò)展的 Vision transformer 模式體系結(jié)構(gòu)，該體系結(jié)構(gòu)僅使用 ImageNet-1K 數(shù)據(jù)集就大大縮小了差距。

VOLO 引入了一種新穎的 outlook attention ，并提出了一種簡(jiǎn)單而通用的架構(gòu)，稱為 Vision Outlooker 。與自我關(guān)注不同，自我關(guān)注側(cè)重于粗略級(jí)別的全局依賴關(guān)系建模， outlook 關(guān)注有效地將更精細(xì)級(jí)別的功能和上下文編碼為標(biāo)記。這對(duì)識(shí)別性能極為有利，但在很大程度上被自我注意所忽視。

實(shí)驗(yàn)表明， VOLO 在 ImageNet-1K 分類上達(dá)到了 87.1% 的 top-1 精度，這是第一個(gè)在這個(gè)競(jìng)爭(zhēng)基準(zhǔn)上超過(guò) 87% 精度的模型，無(wú)需使用任何額外的訓(xùn)練數(shù)據(jù)。

圖 1 ：不同尺寸級(jí)別的 VOLO 模型的 Top-1 精度

此外，經(jīng)過(guò)預(yù)訓(xùn)練的 VOLO 可以很好地轉(zhuǎn)移到下游任務(wù)，例如語(yǔ)義切分。

表 1 ：對(duì)比 ViT 、 CNN 基線模型概述

盡管 VOLO 模型顯示出了出色的計(jì)算效率，但訓(xùn)練 SOTA 性能模型并非易事。

在這篇文章中，我們將介紹我們?cè)?NVIDIA DGX SuperPOD 上基于 NVIDIA ML 軟件堆棧和 Infiniband 群集技術(shù)培訓(xùn) VOLO 模型所獲得的技術(shù)和經(jīng)驗(yàn)。

培訓(xùn)方法

培訓(xùn) VOLO 模型需要考慮培訓(xùn)策略、基礎(chǔ)設(shè)施和配置規(guī)劃。在本節(jié)中，我們將討論此解決方案中應(yīng)用的一些技術(shù)。

培訓(xùn)策略

始終使用原始 ImageNet 樣本質(zhì)量數(shù)據(jù)訓(xùn)練模型，并在細(xì)粒度上執(zhí)行神經(jīng)網(wǎng)絡(luò)（ NN ）架構(gòu)搜索，使理論上的研究更加鞏固。然而，這需要計(jì)算資源預(yù)算的很大一部分。

在這個(gè)項(xiàng)目的范圍內(nèi)，我們采用了一種粗粒度的訓(xùn)練方法，它不像細(xì)粒度的方法那樣能夠訪問(wèn)盡可能多的神經(jīng)網(wǎng)絡(luò)體系結(jié)構(gòu)。然而，它能夠以更少的時(shí)間和更低的資源預(yù)算顯示 EIOF 。在這種替代策略中，我們首先使用分辨率較低的圖像樣本訓(xùn)練潛在的神經(jīng)網(wǎng)絡(luò)候選，然后使用高分辨率圖像進(jìn)行微調(diào)。

在早期的工作中，這種方法在降低邊際模型性能損失的計(jì)算成本方面被證明是有效的。

基礎(chǔ)設(shè)施

實(shí)際上，我們?cè)诒敬闻嘤?xùn)中使用了兩種類型的集群：

一個(gè)用于基礎(chǔ)模型預(yù)訓(xùn)練，它是一個(gè)基于 NVIDIA DGX A100 的 DGX 吊艙，由使用 NVIDIA Mellanox HDR Infiniband 網(wǎng)絡(luò)集群的 5 個(gè) NVIDIA DGX A100 系統(tǒng)組成。

一個(gè)用于微調(diào)，即 NVIDIA DGX SuperPOD ，由 DGX A100 系統(tǒng)和 NVIDIA Mellanox HDR Infiniband 網(wǎng)絡(luò)組成。

圖 2 ：本項(xiàng)目使用的基于 NVIDIA 技術(shù)的軟件棧

軟件基礎(chǔ)設(shè)施在這一過(guò)程中也發(fā)揮了重要作用。圖 2 顯示，除了基礎(chǔ)的標(biāo)準(zhǔn)深度學(xué)習(xí)優(yōu)化 CUDA 庫(kù)（如 cuDNN 和 cuBLAS ）外，我們還廣泛利用 NCCL 、 enroot 、 PyXis 、 APEX 和 DALI 來(lái)實(shí)現(xiàn)培訓(xùn)性能的亞線性可擴(kuò)展性。

DGX A100 POD 集群主要用于使用較小尺寸圖像樣本的基礎(chǔ)模型預(yù)訓(xùn)練。這是因?yàn)榛灸Ｐ皖A(yù)訓(xùn)練的內(nèi)存限制較少，可以利用 NVIDIA A100 GPU 的計(jì)算能力優(yōu)勢(shì)。

相比之下，微調(diào)是在 NVIDIA DGX-2 的 NVIDIA DGX SuperPOD 上執(zhí)行的，因?yàn)槲⒄{(diào)過(guò)程使用更大的圖像，每臺(tái)計(jì)算能力需要更多的內(nèi)存。

培訓(xùn)配置

需要引入句子

表 2 ：模型設(shè)置（對(duì)于所有模型，批大小設(shè)置為 1024 ）

我們?cè)?ImageNet 數(shù)據(jù)集上評(píng)估了我們提出的 VOLO 模型。在培訓(xùn)期間，沒(méi)有使用額外的培訓(xùn)數(shù)據(jù)。我們的代碼基于 PyTorch 、令牌標(biāo)記工具箱和 PyTorch 圖像模型（ timm ）。我們使用帶有標(biāo)記的 LV-ViT-S 模型作為基線。

安裝說(shuō)明

我們使用了 AdamW 優(yōu)化器和線性學(xué)習(xí)率縮放策略 LR = LR基礎(chǔ)x Batch \ u 大小/ 1024 和 5 × 10 ? 2 先前工作建議的重量衰減率，表 3 中給出了所有 VOLO 模型的 LRbase 。

使用隨機(jī)深度。

我們?cè)?ImageNet 數(shù)據(jù)集上訓(xùn)練了 300 個(gè)時(shí)代的模型。

對(duì)于數(shù)據(jù)擴(kuò)充方法，我們使用 CutOut 、 RandAug 和 MixToken 的標(biāo)記目標(biāo)。

我們沒(méi)有使用 MixUp 或 CutMix ，因?yàn)樗鼈兣c MixToken 沖突。

訓(xùn)練前

在本節(jié)中，我們以 VOLO-D5 為例來(lái)演示如何訓(xùn)練模型。

圖 3 顯示，使用單個(gè) DGX A100 的 VOLO-D5 的訓(xùn)練吞吐量約為 500 圖像/秒。據(jù)估計(jì)，完成一個(gè)完整的預(yù)訓(xùn)練周期大約需要 170 個(gè)小時(shí)，這需要使用 ImageNet-1K 進(jìn)行 300 個(gè)階段。這相當(dāng)于 100 萬(wàn)張圖片的一周。

為了加快速度，基于一個(gè)由五個(gè) DGX A100 節(jié)點(diǎn)組成的簡(jiǎn)單參數(shù)服務(wù)器架構(gòu)集群，我們大致實(shí)現(xiàn)了 2100 個(gè)圖像/秒的吞吐量，這可以將預(yù)訓(xùn)練時(shí)間減少到約 52 小時(shí)。

圖 3 ：D1 ～ D5 模型在一個(gè) DGX A100 上跨一個(gè)完整歷元的訓(xùn)練吞吐量

VOLO-D5 模型預(yù)訓(xùn)練可以使用以下代碼示例在單個(gè)節(jié)點(diǎn)上啟動(dòng)：

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

對(duì)于 MNMG 培訓(xùn)案例，它需要將培訓(xùn)集群詳細(xì)信息作為命令行輸入的一部分。首先，我們根據(jù)節(jié)點(diǎn)和集群架構(gòu)設(shè)置 CPU 、 MEM 、 IB 綁定。預(yù)訓(xùn)練階段的集群是 DGX A100 POD ，每個(gè) CPU 插槽有四個(gè) NUMA 域，每個(gè) A100 GPU 有一個(gè) IB 端口，因此我們將每個(gè)列組綁定到 NUMA 節(jié)點(diǎn)中距離其 GPU 最近的所有 CPU 核。

對(duì)于內(nèi)存綁定，我們將每個(gè)列組綁定到最近的 NUMA 節(jié)點(diǎn)。

對(duì)于 IB 綁定，我們?yōu)槊總€(gè) GPU 綁定一個(gè) IB 卡，或者盡可能接近這樣的設(shè)置。

由于 VOLO 模型培訓(xùn)基于 PyTorch ，并且簡(jiǎn)單地利用了默認(rèn)的 PyTorch 分布式培訓(xùn)方法，因此我們的多節(jié)點(diǎn)多 GPU 培訓(xùn)基于一個(gè)簡(jiǎn)單的參數(shù)服務(wù)器架構(gòu)，該架構(gòu)適合 NVIDIA DGX SuperPOD 的 fat 樹(shù)網(wǎng)絡(luò)拓?fù)洹?/p>

為了簡(jiǎn)化調(diào)度，分配節(jié)點(diǎn)列表中的第一個(gè)節(jié)點(diǎn)始終用作參數(shù)服務(wù)器和工作節(jié)點(diǎn)，而所有其他節(jié)點(diǎn)都是工作節(jié)點(diǎn)。為了避免潛在的存儲(chǔ) I / O 開(kāi)銷，數(shù)據(jù)集、所有代碼、中間/里程碑檢查點(diǎn)和結(jié)果都保存在一個(gè)基于 DDN 的高性能分布式存儲(chǔ)后端。它們通過(guò) 100G NVIDIA Mellanox EDR Infiniband 網(wǎng)絡(luò)裝載到所有工作節(jié)點(diǎn)。

為了加速數(shù)據(jù)預(yù)處理和流水線數(shù)據(jù)加載， NVIDIA DALI 配置為每個(gè) GPU 進(jìn)程使用一個(gè)專用數(shù)據(jù)加載程序。

圖 4 ：訓(xùn)練前階段訓(xùn)練吞吐量相對(duì)于 A100 和 V100 的速度提高 GPU

微調(diào)

使用以下代碼示例，在單個(gè)節(jié)點(diǎn)上運(yùn)行 VOLO-D5 模型微調(diào)非常簡(jiǎn)單：

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/

如前所述，由于用于微調(diào)的圖像大小遠(yuǎn)遠(yuǎn)大于預(yù)訓(xùn)練階段使用的圖像大小，因此必須相應(yīng)地減小批量大小。將工作負(fù)載放入 GPU 內(nèi)存中，這使得進(jìn)一步擴(kuò)展訓(xùn)練到更大數(shù)量的 GPU 并行任務(wù)是必須的。

圖 5 ：針對(duì) A100 和 V100 的數(shù)量提高微調(diào)階段訓(xùn)練吞吐量 GPU

大多數(shù)微調(diào)配置類似于預(yù)訓(xùn)練階段。

結(jié)論

在這篇文章中，我們展示了在大規(guī)模人工智能超級(jí)計(jì)算機(jī)上訓(xùn)練 SOTA 大規(guī)模視覺(jué) transformer 模型（如 VOLO \ u D5 ）的主要技術(shù)和程序，如基于 NVIDIA DGX A100 的 DGX SuperPOD 。經(jīng)過(guò)訓(xùn)練的 VOLO \ u D5 模型在圖像分類模型排名中取得了最佳的 Top-1 精度，無(wú)需使用 ImageNet-1k 數(shù)據(jù)集以外的任何其他數(shù)據(jù)。

這項(xiàng)工作的代碼資源（包括用于運(yùn)行實(shí)驗(yàn)的 Docker 映像和 Slurm 調(diào)度程序腳本）在 sail-sg/volo GitHub repo 中是開(kāi)源的，以便將來(lái)可以在 VOLO \ u D5 上進(jìn)行更廣泛的研究。

未來(lái)，我們希望進(jìn)一步擴(kuò)展這項(xiàng)工作，以培訓(xùn)更智能、自我監(jiān)督、更大規(guī)模的模型，使用更大的公共數(shù)據(jù)集和更現(xiàn)代化的基礎(chǔ)設(shè)施，例如， NVIDIA DGX SuperPOD 和 NVIDIA H100 GPU。

關(guān)于作者

Terry Yin 目前是 NVIDIA AI 技術(shù)中心的高級(jí)深度學(xué)習(xí)解決方案架構(gòu)師。他分別于 2009 年和 2012 年在中國(guó)華南理工大學(xué)和韓國(guó)延世大學(xué)獲得學(xué)士和碩士學(xué)位。 2012 年至 2016 年，他是南洋理工大學(xué)新加坡分校的研究員，期間他獲得了東盟 ICT 金獎(jiǎng)、數(shù)據(jù)中心動(dòng)態(tài)獎(jiǎng)、 ACM SIGCOMM 2013 年旅游獎(jiǎng)和 GTC 2015 年演講者獎(jiǎng)。他的研究興趣包括云計(jì)算系統(tǒng)、深度學(xué)習(xí)系統(tǒng)、高性能計(jì)算系統(tǒng)等。

Yuan Lin 是 NVIDIA 編譯團(tuán)隊(duì)的首席工程師。他對(duì)所有使程序更高效、編程更高效的技術(shù)感興趣。在加入 NVIDIA 之前，他是 Sun Microsystems 的一名高級(jí)職員工程師。

審核編輯：郭婷