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

發(fā)布人：計算機(jī)視覺工坊時間：2022-06-17 來源：工程師

加入技術(shù)交流群
- 掃碼加入
  和技術(shù)大咖面對面交流
  海量資料庫查詢

發(fā)布文章

作者丨zzk譯

來源丨GiantPandaCV

近期MetaAI發(fā)布了一篇博客，關(guān)于如何顯著提升Vision Transformer的訓(xùn)練效率。

原文：[Significantly faster Vision Transformer training]

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

What the research is

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

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

How it works ？

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

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

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ù)進(jìn)行修改。而當(dāng)參數(shù)比較多的情況下，對應(yīng)啟動的Optimizer Kernel就會變得很多，通常這些Kernel都比較小，計算負(fù)擔(dān)不大，啟動Kernel的開銷反而占了大頭。

在ContiguousParams中，它將模型參數(shù)放置到一塊連續(xù)的顯存中進(jìn)行計算，這樣就能減少優(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下，達(dá)到了1.51倍的加速比，進(jìn)一步的我們將batchsize設(shè)置為384，并將圖片大小增大到256，達(dá)到了1.86倍加速比。在全FP16運(yùn)算下，能夠達(dá)到2.18倍加速比，盡管這偶爾會降低準(zhǔn)確率（在實驗中，準(zhǔn)確率降低不到10%）。