原作 Felipe Petroski Such, Kenneth O. Stanley and Jeff Clune
Root 编译自 Uber Engineering blog量子位 出品 |
一台足矣。
昨天,优步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