谷歌 人工智能部門 最近開源了 GPipe ，這是一個(gè)用于快速訓(xùn)練大規(guī)模深度學(xué)習(xí)模型的 TensorFlow 類庫(kù)。

　　深層神經(jīng)網(wǎng)絡(luò)（DNN）主要用于解決自然語言處理和視覺目標(biāo)識(shí)別等人工智能任務(wù)。以視覺識(shí)別為例，該領(lǐng)域的最新方法通常以 ImageNet 挑戰(zhàn)賽 的獲勝方案為基準(zhǔn)。每一屆冠軍的成績(jī)都優(yōu)于前一屆；當(dāng)然，模型的復(fù)雜度也會(huì)相應(yīng)增加。2014 年的冠軍 GoogLeNet 通過使用 400 萬個(gè)模型參數(shù)達(dá)到了 74.8% 的 top-1 準(zhǔn)確率，而 2017 年的冠軍 Squeeze-and-Excitation Networks 則使用了 1.458 億個(gè)參數(shù)并達(dá)到了 82.7% 的 top-1 準(zhǔn)確率。

　　在訓(xùn)練神經(jīng)網(wǎng)絡(luò)的時(shí)候，模型大小的增加通常會(huì)引起問題。為了在合理的時(shí)間內(nèi)完成訓(xùn)練，我們把大部分的計(jì)算任務(wù)委托給了加速器：諸如 GPU 和 TPU 之類的專用硬件。但是這些設(shè)備的內(nèi)存有限，這也就限制了訓(xùn)練模型的大小。我們可以通過一些方法來減少模型對(duì)內(nèi)存的依賴，比如將數(shù)據(jù)從加速器內(nèi)存中置換出去，但這會(huì)大大減慢訓(xùn)練速度。另一種解決方案則是模型分區(qū)，這可以讓模型同時(shí)在多個(gè)加速器中并行執(zhí)行。對(duì)順序性 DNN 來說，最好的策略是按層劃分模型，然后由不同的加速器來訓(xùn)練不同的層。但是由于 DNN 的順序性本質(zhì)，有些時(shí)候可能只有一個(gè)加速器在工作，別的加速器則因?yàn)樾枰却渌鼘拥挠?xùn)練結(jié)果而閑置下來。

　　GPipe 通過進(jìn)一步細(xì)化訓(xùn)練任務(wù)解決了這個(gè)問題，它將批量任務(wù)分解為更細(xì)小的“微批量”任務(wù)，并在每一層中管道化執(zhí)行這些“微批量”任務(wù)。這樣，下一層的加速器就可以優(yōu)先處理上一層已完成的“微批量”任務(wù)結(jié)果，而不需要等待整個(gè)訓(xùn)練過程的結(jié)束。

　　通過使用 GPipe 以及 8 個(gè) TPUv2（第二代 TPU 芯片），谷歌研究人員能夠用 18 億個(gè)參數(shù)來訓(xùn)練視覺目標(biāo)識(shí)別模型：在使用 GPipe 的情況下，單個(gè) TPUv2 可訓(xùn)練的參數(shù)量增加了 5.6 倍。通過此次訓(xùn)練的大規(guī)模模型，ImageNet 數(shù)據(jù)驗(yàn)證的準(zhǔn)確率達(dá)到了 84.7%，超過了 2017 年奪冠時(shí)的 82.7%。

　　GPipe 的模型分區(qū)除了能支持更大的模型以外，它也允許多個(gè)加速器并行訓(xùn)練所指定的模型。研究報(bào)告稱，使用 4 倍以上的加速器可以達(dá)到 3.5 倍的加速效果。

　　Gpipe 目前是 Lingvo 框架 的一部分，該框架主要用來在 TensorFlow 中構(gòu)建順序神經(jīng)網(wǎng)絡(luò)模型。