如何更快地訓練Vision Transformer

來源丨GiantPandaCV

近期MetaAI發布了一篇博客，關于如何顯著提升Vision Transformer的訓練效率。

原文：[Significantly faster Vision Transformer training]

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

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

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

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

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

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

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

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

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

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

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

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

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

這里的數據格式與以往不同，將batch維度放在第二維，并在調用nn.MultiheadAttention的時候，設置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

總感覺這個實現怪怪的

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

使用上述優化，我們將Imagenet1K數據集每epoch訓練時間從0.65小時降低到0.43小時

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

本文僅做學術分享，如有侵權，請聯系刪文。