15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

選自arXiv

作者:Ji Lin, Chuang Gan, Song Han機器之心編譯參與:Panda大規模視頻識別分佈式訓練往往具有較高的計算、數據加載和通信成本,為了解決這些問題,MIT 的研究者提出了三項設計原則,並基於這些原則設計了一種時間位移模塊(TSM)。這一設計成功將在 Summit 超級計算機上需要近 50 小時的訓練時間壓縮到了不到 15 分鐘。

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

  • 鏈接:https://arxiv.org/abs/1910.00932
  • TSM Github 鏈接: https://github.com/mit-han-lab/temporal-shift-module
  • 項目地址: https://hanlab.mit.edu/projects/tsm/

深度視頻識別的計算成本比圖像識別更高,尤其是在 Kinetics 等大規模數據集上。因此,為了處理大量視頻,可擴展性訓練是至關重要的。這篇論文研究了影響視頻網絡的可擴展性的因素。研究者認定了三個瓶頸,包括數據加載(從磁盤向 GPU 移動數據)、通信(在網絡中移動數據)和計算速度(FLOPs)。

針對這些瓶頸,研究者提出了三種可以提升可擴展性的設計原則:(1)使用 FLOPs 更低且對硬件友好的算子來提升計算效率;(2)降低輸入幀數以減少數據移動和提升數據加載效率,(3)減少模型大小以降低網絡流量和提升網絡效率。

基於這些原則,研究者設計了一種新型的算子「時間位移模塊(TSM:Temporal Shift Module)」,能夠實現高效且可擴展的分佈式訓練。相比於之前的 I3D 模型,TSM 模型的吞吐量可以高出 1.8 倍。

研究者也通過實驗測試了新提出的 TSM 模型。將 TSM 模型的訓練擴展到了 1536 個 GPU 上,使用了包含 12288 個視頻片段/ 98304 張圖像的 minibatch,沒有造成準確度損失。使用這樣的硬件友好的模型設計,研究者成功地擴展了在 Summit 超級計算機上的訓練,將在 Kinetics 數據集上的訓練時間從 49 小時 55 分減少到了 14 分 13 秒,同時實現了 74.0% 的 top-1 準確度,這在準確度更高的同時還比之前的 I3D 視頻模型快 1.6 和 2.9 倍。

研究介紹

在計算機視覺領域,視頻識別是一個至關重要的分支。視頻識別問題的難度更高,但得到的研究更少:(1)相比於 2D 圖像模型,視頻模型的計算成本通常高一個數量級。舉個例子,很常見的 ResNet-50 模型的速度大約是 4G FLOPs,而 ResNet-50 I3D 則要消耗 33G FLOPs,多過 8 倍;(2)視頻數據集比 2D 圖像數據集大得多,而且數據 I/O 也比圖像高很多。舉個例子,ImageNet 有 128 萬張訓練圖像,而視頻數據集 Kinetics-400 有 6300 萬訓練幀,大約是前者的 50 倍;(3)視頻模型的模型大小通常更大,因此需要更高的網絡帶寬來交換梯度。

這篇論文研究了視頻的大規模分佈式訓練的瓶頸,包括計算、數據加載(I/O)、通信。

對應這些瓶頸,研究者又提出了用於解決這些難題的三項實用的設計原則:模型應當利用對硬件友好的算子來降低計算 FLOPs;模型應當使用更少的輸入幀以節省文件系統 I/O;模型應當使用參數更少的算子以節省網絡帶寬。

基於這些原則,研究者提出了一種零 FLOPs 和零參數的高效視頻 CNN 算子「時間位移模塊(TSM:Temporal Shift Module)」。它可將 Kinetics 訓練擴展到 1536 個 GPU 上,實現包含 12288 個視頻片段/ 98304 張圖像的 minibatch。整個訓練過程可在 15 分鐘內完成,並能實現 74.0% 的 top-1 準確度。相比於之前的兩種 I3D 模型,TSM 模型在當前引領世界的 Summit 超級計算機上可實現分別高 1.6 和 2.9 倍的訓練吞吐量。

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

圖 1:時間位移模塊(TSM)會沿時間維度位移通道,從而實現近鄰幀之間的時間建模

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

圖 2:兩種類型的視頻主幹設計。(a) 沒有時序上的池化,對數據的利用更加高效。(b) 需要更多的輸入幀數,對 I/O 造成了壓力。

視頻模型架構與模型設計原則

為了解決分佈式訓練系統中的難題,研究者提出了三項視頻模型設計原則:(1)為了提升計算效率,使用 FLOPs 更低且硬件效率更高的算子;(2)為了降低數據加載流量,使用「FLOPs/數據」比更高的網絡拓撲結構;(3)為了降低網絡流量,使用參數更少的算子。

藉助於硬件友好的模型設計技術,研究者可將訓練擴展到 1536 個 GPU 上,並能在 15 分鐘內結束 Kinetics 訓練。

實驗設置

在實驗中,研究者使用了目前世界上最快的超級計算機 Summit。其由大約 4600 個計算節點構成,每一個節點都有兩個 IBM POWER9 處理器和六個英偉達 Volta V100 加速器。POWER9 處理器通過兩個 NVLINK bricks 相連,每一個在每個方向都有 25GB/s 的傳輸速度。節點包含供 POWER9 處理器使用的 512GB 內存和供加速器使用的 96GB 高帶寬內存(HBM2)。

分佈式訓練使用了 PyTorch 和 Horovod。該框架使用 ring-allreduce 算法來執行同步隨機梯度下降。訓練使用了 CUDA 和 cuDNN 加速。大多數通信都使用了 NVIDIA Collective Communication Library (NCCL) 2。

實驗使用的數據集為 Kinetics-400:https://deepmind.com/research/open-source/kinetics。該數據集包含 400 個人類動作類別,每個類別包含至少 400 段視頻。該數據集包含大約 24 萬段訓練視頻和 2 萬段驗證視頻,每段視頻持續時間為大約 10 秒鐘。這樣大的規模對模型訓練和數據存儲構成了嚴峻挑戰。

訓練過程持續了 100 epoch。研究者訓練了一個有 8 幀輸入的 TSM 網絡,使用了固定的 n=8。初始學習率設置為每 8 個樣本 0.00125,研究者使用了線性縮放規則來增大批大小更大時的學習率。訓練過程使用了餘弦式的學習率衰減,並有 5 epoch 的預熱。權重衰減為 1e-4,沒有使用 dropout。另外,沒有在批歸一化和偏置上使用權重衰減。

測試方法是每段視頻採樣 10 個片段並計算平均預測結果。視頻的大小都經過了部分調整,短邊大小均調整為 256,然後再輸入網絡。

實驗結果

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

表 1:不同模型的效率統計。箭頭表示越大或越小的結果更好

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

圖 3:分析不同的設計方面將如何影響視頻識別模型的分佈式訓練可擴展性:(a)計算效率;(b)數據加載效率;(c)網絡效率

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

表 2:每個塊的輸出特徵圖的時間分辨率。TSM 是全 2D 結構,具有最佳的硬件效率。I3D3×3×3 的最後幾個階段的時間分辨率更低,這使其更接近 2D CNN,因此相比於 I3D3×1×1 有更好的硬件效率。

15分鐘完成Kinetics視頻識別訓練,除了超級計算機你還需要TSM

圖 6:分佈式同步 SGD 訓練的吞吐量和可擴展性。甚至當使用 1536 個 GPU 時,TSM8f 也能實現很好的可擴展性(>80%)。TSM8f 的訓練吞吐量比 I3D3×3×3 高 1.6 倍,比 I3D3×1×1 高 2.9 倍,這表明新提出的設計原則是有效的。


分享到:


相關文章: