4小時學會雅達利遊戲，AI需要幾台電腦？
原作 Felipe Petroski Such, Kenneth O. Stanley and Jeff Clune
Root 編譯自 Uber Engineering blog量子位 出品 | 公眾號 QbitAI
一台足矣。
昨天，優步AI Lab開源了深度神經進化的加速代碼。其博客上稱，哪怕用戶只有一台電腦（台式機），用這個代碼也能訓練出會打雅達利的AI。而且只需要4！小！時！
要知道，此前用深度神經進化方法，讓AI一小時學會玩雅達利，需要720個CPU。
720個CPU啊……誰來算算要多少錢……
別費勁算了。反正就是貴到爆炸。想訓練出自己的AI？不是壕根本不要想。
優步還是決心伸出援手，拉貧苦大眾一把。他們研發出這個深度神經進化加速代碼，是希望降低AI研究的資金門檻。至少讓那些窮得只能買得起遊戲配置的學生們，想自己動手玩玩AI時還有機會。
用神經進化技術，解決深度強化學習問題是非常promising的。玩雅達利（↖）和MuJoCo機器人運動模擬（↗）都可以用神經進化技術。
只用一台電腦，4小時？怎麼做到的？
弱弱地先補充個小前提：這台式電腦的配置得高端一點。
事實上，高端的台式電腦是有幾十個虛擬核的，這相當於中型計算集群了。
如果能恰當的平行運行評估環節的話，那原本需要720核CPU跑1小時的任務現在48核的個人電腦只要16小時。
但別忘了高端的電腦還有GPU可以用，畢竟GPU跑深度神經網路更快。
- 整體思路：CPU和GPU同時用起來
優步這次開源的代碼可以同時最大化「榨」干CPU和GPU。讓GPU跑深度神經網路，讓CPU跑電子遊戲或物理模擬器，然後平行批量運行多個評估程序。另外，它也包含調整過的TensorFlow運算，這對於訓練速度的提高相當重要。
- 先改進GPU
在優步的setup里，運行單個神經網路，用CPU比GPU要快。
但GPU更擅長平行運行多個神經網路。
所以，為了充分發揮GPU的優勢，優步把多個神經網路整合到了一起，放在同一批里運行。這種做法，在神經網路研究里還是挺常見的，通常用於同樣的神經網路處理不同的輸入的情況。
不過，神經網路的進化是不一樣的。它需要用到不同類的神經網路。這對記憶的容量要求就變高了。
一開始，優步AI研究團隊用基礎的TensorFlow運算跑了一次，只用8小時就搞定。但他們覺得還有空間優化。
於是，他們進一步調整TensorFlow，增加了兩類自定義的運算，於是速度又提升了一倍，訓練用時再降到4小時。
其中，第一個自定義TensorFlow運算顯著加速了GPU。它專門針對強化學習中的異構神經網路強化學習運算，比如雅達利遊戲和大多數機器人運動模擬任務，在這些任務中，各個神經網路的訓練周期長度都不同。有了這個自定義運算，GPU可以只調用必需的神經網路進行運算，而不必每一次任務都跑完所有的網路，這樣能夠大大節省計算量。
- 再針對CPU
上述提到的所有改進，側重點都在GPU上。不過，GPU快到一定程度，也會受制於CPU的速度。為了提高CPU運行遊戲模擬器的表現，就需要第二個自定義的TensorFlow運算。
具體來說，就是把雅達利模擬器的wrapper從Python改成自定義TensorFlow指令（reset，step，observation）。
這樣不僅能利用TensorFlow多線程的快速處理能力，還避免了因為Python和TensorFlow相互作用而造成的減速。
總體上，這些調整極大地提升了雅達利模擬器的運行速度。確切來說，大概是原來的3倍。
對於任何一個域（比如雅達利模擬器、物理模擬器等）里包含多個實例的強化學習研究，上面提到的這些創新方法都能幫它們加速。這類研究越來越常見，比如深度Q學習DQN，或策略梯度A3C的強化學習都是這樣。
分開優化後，整體再協調
單拎GPU或CPU來看，速度上已經提升到極致了。但下一個要面臨的挑戰是如何讓所有的計算資源都同時開動起來。
比方說，我們在每個神經網路上都加個前饋傳遞，問它當下狀態應該採取什麼行動，那麼當每個神經網路計算答案的時候，運行遊戲模擬器的CPU只能空轉。
同樣，如果我們問域模擬器——現在這些行為會產生什麼狀態，那麼CPU在模擬出結果的時候，GPU就沒事幹了。
下圖就是CPU+GPU多線程的示意圖。
從這張圖可以了解到，如何逐步優化強化學習里處理不同類型神經網路的運行效果。
藍框指的是域模擬器，在這篇文章里是雅達利遊戲模擬器，或者是MuJoCo物理模擬器。這兩個模擬器都有不同時間長度的場景。
最左側，是最原始的處理方法，效果最差。既沒有充分利用到平行計算的能力，也沒浪費了GPU或CPU等待對方處理數據的時間。
中間是多線程方法稍微好一點，一個CPU可以同時處理多個模擬任務，等到GPU運行時，也可以接上CPU處理好的多組數據。不過還是存在互相等造成的時間浪費。
優步的方法是CPU+GPU流水線法。其中，CPU不帶停的。在GPU處理CPU數據的時候，CPU並沒有閑下來，繼續馬不停蹄地處理更多的數據。
4百萬參數的神經網路，用優步個方法，4個小時就能搞定。
流水線法的意義
又快又便宜。
這直接降低了研究門檻，使得更多自學AI的人，尤其是學生群體，也可以訓練出自己想要的深度神經網路了。
效率大大提升的代碼將會促進研究更高速的發展。
優步自己就受益匪淺。
他們針對遺傳演算法發布的大型超參數搜索，所花預算是只是原來的一部分。而且在大部分的雅達利遊戲里的表現都獲得了提升。具體數字請看優步更新的論文。（https://arxiv.org/abs/1712.06567）
另外，這代碼還能縮短神經進化迭代的周期，AI研究人員可以更快地嘗試新的想法，也願意去試原先很大耗時很久的網路了。
優步AI lab研發的軟體庫包含深度遺傳演算法的實施，源自Saliman等工作的進化策略演算法，還有優步自家據說超級好用隨機搜索控制。
深度神經進化領域裡，近期還有很多里程碑意義的工作。感興趣的可以進一步了解以下機構的工作：
OpenAI：https://blog.openai.com/evolution-strategies/
DeepMind：https://deepmind.com/blog/population-based-training-neural-networks/
Google Brain：https://research.googleblog.com/2018/03/using-evolutionary-automl-to-discover.html
Sentient：https://www.sentient.ai/blog/evolution-is-the-new-deep-learning/
https://github.com/uber-common/deep-neuroevolution/tree/master/gpu_implementation
— 完 —
歡迎大家關注我們的專欄：量子位 - 知乎專欄
誠摯招聘
量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。
量子位 QbitAI · 頭條號簽約作者
?? ? 追蹤AI技術和產品新動態
推薦閱讀：