深度神經進化加速：只需 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 科技評論下面對其中的改進做個詳細的介紹。

實際上，擁有多個虛擬內核的現代高端桌面 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 小時。

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

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

Uber AI Lab 的新軟體庫包括深度遺傳演算法的實現、來自 Salimas 等人的進化策略演算法，以及（非常具有競爭力的！）隨機搜索控制。他們由衷地希望其他人也使用他們的代碼來加速自己的研究活動。他們也邀請整個研究社區參與構建我們的代碼並改進它，例如，分散式 GPU 訓練和添加為這種類型的計算定製的其他 TensorFlow 運算時，有可能獲得進一步的運算提速。

深度神經進化領域的研究現在非常火熱。除了 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 科技評論編譯整理。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！