如何更快地訓練Vision Transformer

來(lái)源丨GiantPandaCV

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

原文：[Significantly faster Vision Transformer training]

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

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

為了降低門(mén)檻，讓更多人能夠應用ViT，我們開(kāi)發(fā)一系列方法來(lái)加速整個(gè)訓練。我們基于MetaAI的圖像分類(lèi)模型庫PyCls實(shí)現了一系列優(yōu)化，這些優(yōu)化極大的提升了模型訓練過(guò)程的吞吐量：

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

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

為了讓訓練更加高效，我們應用了FSDP訓練策略，他能夠將參數，梯度，優(yōu)化器狀態(tài)分片到各GPU上。在FSDP的幫助下，我們可以用更少的GPU資源構建更大的模型。

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

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

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

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

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

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

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

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

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

總感覺(jué)這個(gè)實(shí)現怪怪的

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

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

我們還研究了不同GPU配置對訓練速度的影響，在不同配置下我們都實(shí)現了比DDP baseline更高的吞吐量。隨著(zhù)GPU增加，吞吐量會(huì )因為設備之間的通信開(kāi)銷(xiāo)略微下降。然而即使在64塊GPU下，我們仍然比DDP基線(xiàn)快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

本文僅做學(xué)術(shù)分享，如有侵權，請聯(lián)系刪文。