谷歌開源 GPipe 庫,主要用於大規模深度學習模型的快速訓練
谷歌人工智慧部門 最近開源了GPipe ,這是一個用於快速訓練大規模深度學習模型的TensorFlow 類庫。
深層神經網路(DNN)主要用於解決自然語言處理和視覺目標識別等人工智慧任務。以視覺識別為例,該領域的最新方法通常以ImageNet 挑戰賽 的獲勝方案為基準。每一屆冠軍的成績都優於前一屆;當然,模型的複雜度也會相應增加。2014 年的冠軍GoogLeNet 通過使用 400 萬個模型引數達到了 74.8% 的 top-1 準確率,而 2017 年的冠軍Squeeze-and-Excitation Networks 則使用了 1.458 億個引數並達到了 82.7% 的 top-1 準確率。
在訓練神經網路的時候,模型大小的增加通常會引起問題。為了在合理的時間內完成訓練,我們把大部分的計算任務委託給了加速器:諸如 GPU 和TPU 之類的專用硬體。但是這些裝置的記憶體有限,這也就限制了訓練模型的大小。我們可以通過一些方法來減少模型對記憶體的依賴,比如將資料從加速器記憶體中置換出去,但這會大大減慢訓練速度。另一種解決方案則是模型分割槽,這可以讓模型同時在多個加速器中並行執行。對順序性 DNN 來說,最好的策略是按層劃分模型,然後由不同的加速器來訓練不同的層。但是由於 DNN 的順序性本質,有些時候可能只有一個加速器在工作,別的加速器則因為需要等待其它層的訓練結果而閒置下來。
GPipe 通過進一步細化訓練任務解決了這個問題,它將批量任務分解為更細小的“微批量”任務,並在每一層中管道化執行這些“微批量”任務。這樣,下一層的加速器就可以優先處理上一層已完成的“微批量”任務結果,而不需要等待整個訓練過程的結束。
通過使用 GPipe 以及 8 個 TPUv2(第二代 TPU 晶片),谷歌研究人員能夠用 18 億個引數來訓練視覺目標識別模型:在使用 GPipe 的情況下,單個 TPUv2 可訓練的引數量增加了 5.6 倍。通過此次訓練的大規模模型,ImageNet 資料驗證的準確率達到了 84.7%,超過了 2017 年奪冠時的 82.7%。
GPipe 的模型分割槽除了能支援更大的模型以外,它也允許多個加速器並行訓練所指定的模型。研究報告稱,使用 4 倍以上的加速器可以達到 3.5 倍的加速效果。
Gpipe 目前是Lingvo 框架 的一部分,該框架主要用來在 TensorFlow 中構建順序神經網路模型。
檢視英文原文:Google Open-Sources GPipe Library for Faster Training of Large Deep-Learning Models