深度神經進化加速:只需 4 個小時就在個人電腦上訓練出強化學習模型

深度神经进化加速:只需 4 个小时就在个人电脑上训练出强化学习模型

雷鋒網 AI 科技評論按:在之前的文章 Uber論文5連發宣告神經演化新時代中,我們介紹了 Uber AI Lab 在深度進化算法方面的研究成果,從多個角度展現了進化算法解決強化學習問題的能力,也展現了進化策略 ES(evolution strategies)與梯度下降之間的聯繫。這些研究成果非常精彩、給人啟迪。不過當時我們沒有提到的是,這些成果消耗了相當多的計算資源:實際上論文中的實驗是在 720 到 3000 個 CPU 組成的大規模高性能計算集群上運行的,這樣的集群固然有充沛的計算能力運行進化算法,但在 Uber AI Lab 的研究人員們看來,這種級別的計算能力要求也就把領域內多數的研究人員、學生、企業以及愛好者攔在了門外。

近日 Uber AI Lab 開源了一組進化算法代碼,它的特點是可以高速(同時也更廉價地)進行進化策略研究。根據介紹,訓練神經網絡玩 Atari 遊戲的時間可以從原來在 720 個 CPU 組成的集群上花費 1 個小時,到現在在一臺桌面級電腦上只需要 4 個小時。這一點很重要,因為它極大地刷新了我們對進行這類研究所需的資源多少的認識,從而使更多的研究人員能夠著手研究。雷鋒網 AI 科技評論下面對其中的改進做個詳細的介紹。

深度神经进化加速:只需 4 个小时就在个人电脑上训练出强化学习模型
深度神经进化加速:只需 4 个小时就在个人电脑上训练出强化学习模型

神經進化技術是解決具有挑戰性的深層強化學習問題的一種有競爭力的替代方案,如玩 Atari 遊戲以及模仿人類運動。圖示出了用簡單遺傳算法訓練的深度神經網絡的行為。

是哪些修改使它更快,且可在單臺計算機上運行?

實際上,擁有多個虛擬內核的現代高端桌面 PC 本身就像一箇中型計算集群。如果能正確地並行執行訓練過程,在 720 個核心上如果需要運行 1 個小時的訓練過程,在帶有 48 核 CPU 的個人計算機上運行就需要 16 個小時。這雖然也是較慢的,但不會令人望而卻步。不過,現代臺式機還有 GPU,它們運行深度神經網絡(DNN)的速度很快。Uber AI Lab 的代碼能夠最大化並行使用 CPU 和 GPU。它在 GPU 上運行深度神經網絡,CPU 上運行要訓練的這個任務(例如電子遊戲或物理仿真器),並可以在每個批當中並行運行多個訓練過程,從而可有效地利用所有可用的硬件。如下所述,它還包含自定義的 TensorFlow 操作,這顯著提高了訓練速度。

允許在 GPU 上進行訓練需要對神經網絡的計算過程進行一些修改。在 Uber AI Lab 的研究人員的設置中,單個神經網絡在單個 CPU 上的速度比在 GPU 上更快,但是 GPU 在大批量類似的並行計算(例如,神經網絡的前饋傳播)時有很大好處。為了儘量榨乾 GPU 的計算能力,他們將多個神經網絡的前饋傳播聚合成批次。這樣做在神經網絡研究中是常見的,但通常是同一個網絡處理不同的輸入。然而,進化算法中上有一批參數不同的神經網絡,但是即使網絡不同,也可以用同樣的做法進行加速(雖然內存的需求會相應增加)。他們用基本的 TensorFlow 運算實現了神經網絡群的批量操作,並且它產生了大約 2 倍的加速,把訓練時間減少到了大約 8 小時。

不僅如此,研究人員們還覺得他們可以做得更好。雖然 TensorFlow 提供了所有需要的運算指令,但這些運算並不那麼適合於這種類型的計算。因此,他們添加了兩種自定義的 TensorFlow 運算,加起來可以再把速度提升兩倍,相比在之前一臺機器上最初提到的 16 小時,將訓練減少到大約 4 小時。

第一種定製的 TensorFlow 運算顯著加快了 GPU 上的運算速度。它是專為異構神經網絡計算定製的,在 RL 領域,每步運算操作所需時間具有不同的長度,這在 Atari 遊戲和許多模擬機器人學習任務中是確實存在的。它允許 GPU 只運行需要運行的那幾個網絡,而不需要在每次迭代中都運行整批固定的網絡集。

到目前為止所描述的改進使得 GPU 比 CPU 更具成本效益。事實上,GPU 是如此之快,以至於運行在 CPU 上的 Atari 模擬器無法跟上,即使已經使用了多處理庫做並行化加速計算。為了提高仿真性能,研究人員們添加了第二組定製的 TensorFlow 運算。這些將 Atari 模擬器的接口從 Python 改為定製的 TensorFlow 命令(reset,step,observation),利用了 TensorFlow 提供的多線程快速處理能力,因而沒有 Python 與 TensorFlow 交互時的典型速度下降問題。

總的來說,所有這些變化使得雅達利模擬器獲得了大約 3 倍加速。這些創新應該可以加快任何有多個並行任務實例的強化學習研究(例如 Atari 遊戲或 MujoCo 物理模擬器),這種多實例的做法在強化學習中也是越來越常見,例如分佈式深度 Q 學習(DQN)和分佈式策略梯度(例如A3C)。

只要有能力在 GPU 上運行多個網絡和及在 CPU 上運行更快的任務模擬器,挑戰就只剩下瞭如何讓計算機上的所有資源儘可能地運行。如果我們對每一個神經網絡做了一個前饋傳播,詢問它在當前狀態下應該採取什麼行動,那麼當每個神經網絡都在計算答案時,運行遊戲模擬器的 CPU 就空閒下來了。同樣,如果我們採取了行動,並問任務模擬器「什麼狀態會從這些行為中產生?」在模擬步驟中,運行神經網絡的 GPU 將空閒(注:下圖的左 1)。如果改為了(下圖的第 2 個)多線程 CPU + GPU 選項,從圖中可以看到雖然改進了單線程計算,但仍然是低效的。

一個更好的解決方案是將兩個或多個神經網絡子集與模擬器配對,並且始終保持 GPU 和 CPU 同時運行,這取決於準備採取哪一個步驟(神經網絡或模擬器)來更新來自不同集合的網絡或模擬器。這種方法是最右邊的「流水線 CPU + GPU」選項,如下圖右 1 所示。通過它以及上文提到的其他改進,使得我們訓練 4 百萬參數的神經網絡所需時間降到在一臺計算機上只需 4 小時。

深度神经进化加速:只需 4 个小时就在个人电脑上训练出强化学习模型

在 RL 中優化異構網絡的種群調度。藍色框是任務模擬器,比如 Atari 模擬器或像 MujoCo 這樣的物理引擎,它可以有不同長度的運行時間片。直白的運行方式雖然使用了 GPU(左)但性能低,原因有兩個:1)GPU 的批大小為 1,沒有利用它的並行計算能力,2)存在 GPU 等待 CPU 的空閒時間,反之亦然。多線程方法(中間)允許通過多個 CPU 並行地運行模擬器,來減少 GPU 的等待時間,但是當 CPU 工作時,GPU 空閒,反之亦然。Uber AI Lab 設計的流水線實現(右)允許 GPU 和 CPU 有效地運行。這種方法也適用於同時運行的多個 GPU 和 CPU,實際上他們也就是這樣做的。

實驗更快、更便宜後的意義

Uber AI Lab 的代碼使研究社區的每一個人,包括學生和自學成才的學生,能夠快速實驗性地反覆訓練諸如玩 Atari 遊戲的挑戰性深度神經網絡,而後者是迄今為止僅限於資金充足的工業和學術實驗室的奢侈品。

更快的代碼會帶來研究進展。例如,新代碼使 Uber AI Lab 能夠只花很少的成本就可以為遺傳算法展開一個廣泛的超參數搜索,與他們最初報告的性能相比,將改進大多數 Atari 遊戲的性能。論文已經發表在 arXiv。同樣地,更快的代碼也催化了研究的進步,通過縮短迭代時間來改進深度神經進化,使他們能夠嘗試更多的任務中的每一個新的想法,並且更長時間地運行算法。

深度神經進化領域的研究現在非常火熱。除了 Uber AI Lab 自己的研究和 OpenAI 的研究,最近也有來自 DeepMind、谷歌大腦和Sentient的深度學習進展。Uber AI Lab 希望通過開源使他們的代碼有助於這個領域的發展。

以及最根本地,Uber AI Lab 的目標是降低進行這項研究的成本,使所有背景的研究者能夠嘗試自己的想法來改進深層神經進化,並利用它來實現他們的目標。

即便文中提到的 48 核 CPU 「桌面 PC」引起了一些爭議(7k 人民幣的 AMD 銳龍 Threadripper 1950X16 核 32 線程,15k 人民幣的 Intel i9 7980XE 18 核 36 線程),但也確實是十分有價值的研究成果。

論文地址:https://arxiv.org/abs/1712.06567

via: Accelerating Deep Neuroevolution: Train Atari in Hours on a Single Personal Computer,雷鋒網 AI 科技評論編譯整理。


分享到:


相關文章: