晓查 发自 凹非寺
量子位 报道 | 公众号 QbitAI
对土豪来说最痛苦的是什么, 就是有一大堆硬件却不能实现 1+1=2 的效果.
AI 训练中的并行计算就是如此, 纵使你有一千张 GPU, 也无法实现单机训练一千倍的效果.
最近, 不差钱的谷歌, 开源了一种 SEED RL 框架, 可以更容易地将 AI 训练放在几千台机器上运行, 效果比之前的方法最高提升近 4 倍 .
如果你也不差钱, 在云端上进行大规模并行计算, 那么可以节约 80% 的训练成本. 考虑到现在一个大型 AI 模型动辄上百万的训练费用, 真的是非常可观了.
在训练 AI 玩足球游戏这件事情上, SEED RL 的处理速度可以达到 每秒 240 万帧 . 如果以 60fps 计算, 相当于每秒可处理 11 小时的游戏画面.
SEED RL 体系架构
上一代的分布式强化学习智能体 IMPALA, 其体系架构中包含 Actor 和 Learner 两部分.
Actor 通常在 CPU 上运行, 并且在环境中采取的步骤与对模型进行推断之间进行迭代, 以预测下一个动作.
Actor 经常会更新推理模型的参数, 并且在收集到足够数量的观测结果后, 会将观测结果和动作的轨迹发送给 Learner, 从而对 Learner 进行优化.
在这种架构中, Learner 使用来自数百台机器上的分布式推理输入在 GPU 上训练模型.
但 IMPALA 存在着许多缺点:
1, 使用 CPU 进行神经网络推理, 效率低下. 而且随着模型变大, 运算量变大, 问题会越来越严重.
2,Actor 和 Learner 之间模型参数的带宽成为性能的瓶颈.
3, 资源利用效率低, Actor 在环境和推理两个任务之间交替进行, 而这两个任务的计算要求不同, 很难在同一台机器上充分利用资源.
SEED RL 体系架构解决了以上这些缺点. Actor 可以在 GPU,TPU 这类 AI 硬件加速器上完成推理, 通过确保将模型参数和状态保持在本地来加快推理速度, 并避免数据传输瓶颈.
与 IMPALA 体系结构相反, SEED RL 中的 Actor 仅在环境中执行操作. Learner 在硬件加速器上使用来自多个 Actor 的成批数据来集中执行推理.
SEED RL 在每个环境步骤将观测结果发送给 Learner 的同时, 使用 gPRC 框架的网络库, 将延迟保持在较低水平. 这使 SEED RL 在一台机器上每秒最多可以实现一百万个查询.
Learner 可以扩展到几千个核心上, Actor 的数量可以扩展到几千台机器, 从而实现每秒百万帧的训练速度.
SEED RL 用到了两种最先进的算法: V-trace 和 R2D2 .
V-trace 负责从采样的动作中预测动作的分布, R2D2 负责根据动作的预测未来值选择一个动作.
V-trace 是基于策略梯度的方法, 最早被 IMPALA 采用. 由于 Actor 和 Learner 是异步执行, 而 V-trace 在异步体系架构中的效果很好
第二种算法是 R2D2, 这是一种 Q 学习方法, DeepMind 曾用这种算法将强化学习智能体在 Atari 游戏上的水平提高了 4 倍, 并在 52 款游戏上超过了人类水平.
这种方法允许 Q 学习算法大规模硬件上运行的同时仍然可以使用 RNN.
实验结果
谷歌在 DeepMind 最近研究的开源足球游戏项目 Google Research Football 中进行了基准测试.
使用 64 个 Cloud TPU 核心实现了每秒 240 万帧的数据传输速度, 与之前的最新分布式 IMPALA 相比, 提高了 80 倍 .
IMPALA 要想达到相同的速度, 需要 14000 个 CPU, 而 SEED RL 只用了 4160 个 CPU. 对于相同的速度, IMPALA 需要的 CPU 是 SEED RL 的 3~4 倍.
通过对并行计算的硬件加速器进行优化, 我们就可以放心大胆地提高模型的大小.
在上面的足球游戏任务中, 通过增加模型大小和输入分辨率, 可以解决从前未能解决的一些困难, 让训练模型的效率大幅提高.
传送门
论文地址: https://arxiv.org/abs/1910.06591
GitHub 地址: https://github.com/google-research/seed_rl
版权所有, 未经授权不得以任何形式转载及使用, 违者必究.
来源: http://www.tuicool.com/articles/YFnYbi7