如何更快地訓(xùn)練Vision Transformer

原文：Significantly faster Vision Transformer training

鏈接：https://ai.facebook.com/blog/significantly-faster-vision-transformer-training

What the research is

Vision Transformer模型幾乎火遍計算機視覺各個領(lǐng)域，其性能隨著參數(shù)增加和更久的訓(xùn)練過程而得到提升。隨著模型越來越大，超過了萬億次浮點運算的規(guī)模，該領(lǐng)域達到了瓶頸：訓(xùn)練一個模型往往要一個月，需要幾百上千個GPU，導(dǎo)致大部分人無法接觸到大規(guī)模ViT模型，并進而增加了對加速器的需求。

為了降低門檻，讓更多人能夠應(yīng)用ViT，我們開發(fā)一系列方法來加速整個訓(xùn)練。我們基于MetaAI的圖像分類模型庫PyCls實現(xiàn)了一系列優(yōu)化，這些優(yōu)化極大的提升了模型訓(xùn)練過程的吞吐量：

How it works ？

我們首先對代碼庫進行分析，以定位訓(xùn)練效率低下的原因，最后關(guān)注點落在計算類型上：大部分模型都是用FP32進行訓(xùn)練，如果使用FP16訓(xùn)練的話，可以降低顯存占用，并提高模型訓(xùn)練速度，但這一做法經(jīng)常會導(dǎo)致準(zhǔn)確率下降。

所以我們選了一個折中的方法：自動混合精度。在該方法下，我們用half類型進行計算，以加快訓(xùn)練，減少顯存使用。并以fp32類型存儲參數(shù)，以保證模型準(zhǔn)確率。其中我們沒有手動將網(wǎng)絡(luò)各部分轉(zhuǎn)換成half類型，而是應(yīng)用AMP各種模式（如O1, O2, O3)，以尋找一個既能提升速度又不影響精度的平衡點。

FSDP

為了讓訓(xùn)練更加高效，我們應(yīng)用了FSDP訓(xùn)練策略，他能夠?qū)?shù)，梯度，優(yōu)化器狀態(tài)分片到各GPU上。在FSDP的幫助下，我們可以用更少的GPU資源構(gòu)建更大的模型。

FSDP策略可以參考 [數(shù)據(jù)并行Deep-dive: 從DP 到 Fully Sharded Data Parallel （FSDP）完全分片數(shù)據(jù)并行] 鏈接：https://zhuanlan.zhihu.com/p/485208899

MTA Optimizer

前向計算完畢后，優(yōu)化器需要對各個參數(shù)進行修改。而當(dāng)參數(shù)比較多的情況下，對應(yīng)啟動的Optimizer Kernel就會變得很多，通常這些Kernel都比較小，計算負擔(dān)不大，啟動Kernel的開銷反而占了大頭。

在ContiguousParams中，它將模型參數(shù)放置到一塊連續(xù)的顯存中進行計算，這樣就能減少優(yōu)化器這部分的時間。下圖是Resnet50+SGD是否應(yīng)用ContiguousParams的比較，可以看到OptimizerStep這部分時間顯著減少了。

而NVIDIA的Apex庫的做法則是在底層重新實現(xiàn)了一系列MultiTensorOptimizer，如Adam, Adagrad等等。

Apex這種方法比較硬核，普通用戶如果想要自己自定義優(yōu)化器并應(yīng)用Multi Tensor的優(yōu)化，就必須改動底層CUDA代碼。而最近PyTorch也在計劃提供了一系列foreach接口[Replace optimizers in torch.optim with the ones from torch.optim._multi_tensor] 鏈接：https://github.com/pytorch/pytorch/pull/49039，讓用戶只需要在Python層即可享受到優(yōu)化，對應(yīng)的MultiTensor版Momentum優(yōu)化器代碼如下所示：

torch._foreach_mul_(bufs,?momentum)
torch._foreach_add_(bufs,?grads,?alpha=1?-?dampening)

Pooled Classifier

原版的ViT是額外加了一個分類token，來輸出最后的分類結(jié)果。而這里采用平均池化，如：https://github.com/facebookresearch/pycls/blob/main/pycls/core/config.py#L205 處理最后的分類

Batch Second Input Tensor Layout

這里的數(shù)據(jù)格式與以往不同，將batch維度放在第二維，并在調(diào)用nn.MultiheadAttention的時候，設(shè)置batch_first=False，以減少不必要的轉(zhuǎn)置

if?self.batch_first?and?is_batched:
????return?attn_output.transpose(1,?0),?attn_output_weights
else:
????return?attn_output,?attn_output_weights

總感覺這個實現(xiàn)怪怪的

其他優(yōu)化

我們在采取560大小的batchsize下，達到了1.51倍的加速比，進一步的我們將batchsize設(shè)置為384，并將圖片大小增大到256，達到了1.86倍加速比。在全FP16運算下，能夠達到2.18倍加速比，盡管這偶爾會降低準(zhǔn)確率（在實驗中，準(zhǔn)確率降低不到10%）。

使用上述優(yōu)化，我們將Imagenet1K數(shù)據(jù)集每epoch訓(xùn)練時間從0.65小時降低到0.43小時

我們還研究了不同GPU配置對訓(xùn)練速度的影響，在不同配置下我們都實現(xiàn)了比DDP baseline更高的吞吐量。隨著GPU增加，吞吐量會因為設(shè)備之間的通信開銷略微下降。然而即使在64塊GPU下，我們?nèi)匀槐菵DP基線快1.83倍

文中鏈接

PyCls ：https://github.com/facebookresearch/pycls

ContiguousParams：https://github.com/PhilJd/contiguous_pytorch_params

Adam：https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_adam.cu

公眾號后臺回復(fù)“CVPR 2022”獲取論文合集打包下載～