2020-11-04 21:12:27

作者 | 青暮

近年來，隨着深度學習走向應用落地，快速訓練ImageNet成爲許多機構競相追逐的目標。

2017年6月，FaceBook使用256塊GPU以1小時在ImageNet上成功訓練了ResNet-50。在這之後，ResNet-50基本成爲了快速訓練ImageNet的標配架構。

2018年7月，騰訊實現用2048塊Tesla P40以6.6分鐘在ImageNet上訓練ResNet-50。同年8月，Fast.ai實現了18分鐘的快速訓練，他們使用的硬件是128塊Tesla V100。到11月，索尼更是將訓練時間壓縮爲224秒，不過其當時使用了2176塊Tesla V100。Fast.ai和索尼使用的網絡架構也都是ResNet。

那麼如今，在快速訓練ImageNet的競逐上，形成了什麼局面呢？

1 ResNet霸榜，偶有對手

在DAWNBench的ImageNet訓練排行榜上，我們可以看到，前五名都是使用了ResNet-50，並且最快訓練時間達到了2分38秒，同時還能實現93.04%的top-5準確率，以及14.42美元的低訓練成本。

DAWNBench 是斯坦福發佈的一套基準測試，主要關注端到端的深度學習訓練和推斷過程，用於量化不同優化策略、模型架構、軟件框架、雲和硬件的訓練時間、訓練成本、推理延遲和推理成本。

排行榜地址：https://dawn.cs.stanford.edu/benchmark/ImageNet/train.html

但除了ResNet-50，就沒有其它適合快速訓練ImageNet的架構了嗎？

在這個排行榜的第10位，我們看到了一個孤獨的名字——AmoebaNet-D N6F256，根據排行榜的數據，它用1/4個TPUv2 Pod和1小時的時間在ImageNet上達到了93.03%的top-5準確率。

實際上，該架構在2018年4月就由谷歌提出，在相同的硬件條件下，AmoebaNet-D N6F256的訓練時間比ResNet-50要短很多。

並且，AmoebaNet-D N6F256的訓練成本也只有ResNet-50的一半。AmoebaNet-D是基於進化策略進行架構搜索的NAS架構，谷歌通過提供複雜的構建模塊和較好的初始條件來參與進化過程，實現了手動設計和進化的有機組合。

論文：Regularized Evolution for Image Classifier Architecture Search

除此之外，在DAWNBench的排行榜上，就看不到其它類型的架構了。但是在Papers With Code的ImageNet排行榜上，就沒有出現ResNet和AmoebaNet-D霸榜的現象。

在top-5準確率前五名中，可以注意到，除了ResNet以外，剩下的架構都和EfficientNet有關。其中排名第一的架構，其top-5準確率已經達到了98.7%。不過這個排行榜並沒有列出訓練時間、硬件和成本，所以和DAWNBench不能一概而論。

排行榜地址：https://www.paperswithcode.com/sota/image-classification-on-imagenet

EfficientNet近日在這項競逐中也出現了新突破，新加坡國立大學的尤洋和谷歌研究院的Quoc Le等人發表了一項新研究，表示其用1小時實現了ImageNet的訓練，並且top-1準確率達到了83%（Papers With Code的排行榜前十名的top-1準確率在86.1%到88.5%之間）。

論文地址：https://arxiv.org/pdf/2011.00071.pdf

尤洋在推特上表示，這項研究在準確率足夠高的前提下，在速度上創造了一個世界記錄。

EfficientNets是基於有效縮放的新型圖像分類卷積神經網絡系列。

目前，EfficientNets的訓練可能需要幾天的時間；例如，在Cloud TPU v2-8節點上訓練EfficientNet-B0模型需要23個小時。

在這項研究中，作者探索了在2048個內核的TPU-v3 Pod上訓練EfficientNets的技術，目的是在以這種規模進行訓練時可以實現加速。

作者討論了將訓練擴展到1024個TPU-v3內核、批量大小爲65536時所需的優化，例如大批量優化器的選擇和學習率的規劃，以及分佈式評估和批量歸一化技術的利用。

此外，作者還提供了在ImageNet數據集上訓練EfficientNet模型的時序和性能基準，以便大規模分析EfficientNets的行爲。

通過優化後，作者能夠在1小時4分鐘內將ImageNet上的EfficientNet訓練到83％的top-1準確率。

圖1展示了對不同TPU內核數，EfficientNet-B2和B5訓練時間達到的峯值準確率。訓練時間在分佈式訓練和評估循環初始化之後立即開始計算，並在模型達到峯值準確率時結束。

可以觀察到EfficientNet-B2在1024個TPU-v3內核上的訓練時間爲18分鐘，top-1準確率爲79.7％，全局批處理大小爲32768，這表明訓練速度明顯提高。

通過在1024個TPU-v3內核上將全局批處理規模擴展到65536，可以在EfficientNet-B5上在1小時4分鐘內達到83.0％的準確率。

表2展示了EfficientNet-B2和B5峯值準確率的基準，這裏可以找到準確率的完整基準以及相應的批次大小，作者表示，即使使用全局批次大小，使用他們的方法，也可以在EfficientNet-B5上保持83％的準確率。

這項研究至少表明了，在EfficientNet上探索大規模圖像分類的快速訓練也是一個有希望的方向。

 

2方法

將EfficientNet培訓擴展到1024個TPU-v3內核會帶來很多挑戰，必須通過算法或系統優化來解決。

第一個挑戰是隨着全局批量規模的增加保持模型質量。由於全局批處理規模隨用於訓練的內核數量而定，因此，必須利用大批量訓練技術來保持準確率。

而在大量TPU芯片上進行訓練時也會面臨計算瓶頸，不過可以使用Kumar等人提出的分佈式評估和批量歸一化技術來解決這些瓶頸。以下是爲擴展TPU-v3 Pod上的EfficientNet訓練而探索的優化技術：

大批量優化器

最初的EfficientNet論文使用RMSProp優化器，但衆所周知，批量較大時，RMSProp會導致模型質量下降。

通過數據並行性將訓練規模擴展到1024個TPU-v3內核，意味着如果保持每個內核的批量大小不變，則全局批量大小必須與內核的數量成比例增加。例如，如果我們將每核的批處理大小固定爲32，則1024個核上的全局批處理大小將爲32768。

另一方面，如果在擴展到多個核時全局批處理大小是固定的，則每核的批處理大小會降低，導致效率低下和吞吐量降低。

因此，爲了更好地利用每個TPU內核的內存並提高吞吐量，必須使用較大的全局批處理大小。使用You等人提出的分層自適應速率縮放（LARS）優化器，能夠擴展到65536的批量，同時達到與Tan和Le中報告的EfficientNet基準準確率相似的準確率。

學習率規劃

爲了在使用大批量時保持模型質量，作者還採用了學習率預熱和線性縮放技術。

在保持epoch數固定的同時增加全局批量大小將導致更新權重的迭代數較少。爲了解決這個問題，作者將線性縮放規則應用於批次中每256個樣本的學習率。

但是，較高的學習率會導致發散。因此，作者還應用了學習率預熱，其中訓練以較小的初始學習率開始，並在可調整的時期內逐漸提高學習率。此外，作者比較了各種學習速率規劃，例如指數衰減和多項式衰減，發現對於LARS優化器，多項式衰減計劃可實現最高的準確率。

分佈式評估

評估循環的執行是EfficientNet的標準雲TPU實現的另一個計算瓶頸，因爲評估和訓練循環是在單獨的TPU上執行的。

在傳統的TPUEstimator中，評估是在單獨的TPU上進行的，訓練的執行速度要比評估快，導致端到端時間很大程度上取決於評估時間。

爲了克服這個問題，作者利用Kumar等人所述的分佈式訓練和評估循環。它在所有TPU上分配了訓練和評估步驟，並允許擴展到更大數量的副本。

分佈式批處理歸一化

作者使用Ying等人提出的方案，通過將副本的子集分組在一起，在副本之間分佈批處理歸一化。這種優化可通過權衡TPU之間的通信成本來提高最終精度。分組在一起的副本數是可調超參數。

最終的批次歸一化批次大小（每個副本子集中的樣本總數）也會影響模型質量以及收斂速度。對於大於16的副本子集，作者還探索了將副本分組在一起的二維切片方法。

精度

目前已經觀察到，使用bfloat16浮點格式來訓練卷積神經網絡可以達到甚至超過使用傳統單精度格式（例如fp32）訓練的網絡的性能，這可能是較低精度的正則化效果所致。

作者實現了混合精度訓練，以利用bfloat16的性能優勢，同時仍保持模型質量。在實驗中，bfloat16用於卷積運算，而所有其他運算都使用fp32。使用bfloat16格式進行卷積可提高硬件效率，而不會降低模型質量。