本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

對土豪來說最痛苦的是什么，就是有一大堆硬件卻不能實(shí)現(xiàn)1+1=2的效果。

AI訓(xùn)練中的并行計(jì)算就是如此，縱使你有一千張GPU，也無法實(shí)現(xiàn)單機(jī)訓(xùn)練一千倍的效果。

最近，不差錢的谷歌，開源了一種SEED RL框架，可以更容易地將AI訓(xùn)練放在幾千臺機(jī)器上運(yùn)行，效果比之前的方法最高提升近4倍。

如果你也不差錢，在云端上進(jìn)行大規(guī)模并行計(jì)算，那么可以節(jié)約80%的訓(xùn)練成本?？紤]到現(xiàn)在一個(gè)大型AI模型動(dòng)輒上百萬的訓(xùn)練費(fèi)用，真的是非?？捎^了。

在訓(xùn)練AI玩足球游戲這件事情上，SEED RL的處理速度可以達(dá)到每秒240萬幀。如果以60fps計(jì)算，相當(dāng)于每秒可處理11小時(shí)的游戲畫面。

SEED RL體系架構(gòu)

上一代的分布式強(qiáng)化學(xué)習(xí)智能體IMPALA，其體系架構(gòu)中包含Actor和Learner兩部分。

Actor通常在CPU上運(yùn)行，并且在環(huán)境中采取的步驟與對模型進(jìn)行推斷之間進(jìn)行迭代，以預(yù)測下一個(gè)動(dòng)作。

Actor經(jīng)常會(huì)更新推理模型的參數(shù)，并且在收集到足夠數(shù)量的觀測結(jié)果后，會(huì)將觀測結(jié)果和動(dòng)作的軌跡發(fā)送給Learner，從而對Learner進(jìn)行優(yōu)化。

在這種架構(gòu)中，Learner使用來自數(shù)百臺機(jī)器上的分布式推理輸入在GPU上訓(xùn)練模型。

但I(xiàn)MPALA存在著許多缺點(diǎn)：

1、使用CPU進(jìn)行神經(jīng)網(wǎng)絡(luò)推理，效率低下。而且隨著模型變大、運(yùn)算量變大，問題會(huì)越來越嚴(yán)重。

2、Actor和Learner之間模型參數(shù)的帶寬成為性能的瓶頸。

3、資源利用效率低，Actor在環(huán)境和推理兩個(gè)任務(wù)之間交替進(jìn)行，而這兩個(gè)任務(wù)的計(jì)算要求不同，很難在同一臺機(jī)器上充分利用資源。

SEED RL體系架構(gòu)解決了以上這些缺點(diǎn)。Actor可以在GPU、TPU這類AI硬件加速器上完成推理，通過確保將模型參數(shù)和狀態(tài)保持在本地來加快推理速度，并避免數(shù)據(jù)傳輸瓶頸。

與IMPALA體系結(jié)構(gòu)相反，SEED RL中的Actor僅在環(huán)境中執(zhí)行操作。Learner在硬件加速器上使用來自多個(gè)Actor的成批數(shù)據(jù)來集中執(zhí)行推理。

SEED RL在每個(gè)環(huán)境步驟將觀測結(jié)果發(fā)送給Learner的同時(shí)，使用gPRC框架的網(wǎng)絡(luò)庫，將延遲保持在較低水平。這使SEED RL在一臺機(jī)器上每秒最多可以實(shí)現(xiàn)一百萬個(gè)查詢。

Learner可以擴(kuò)展到幾千個(gè)核心上，Actor的數(shù)量可以擴(kuò)展到幾千臺機(jī)器，從而實(shí)現(xiàn)每秒百萬幀的訓(xùn)練速度。

SEED RL用到了兩種最先進(jìn)的算法：V-trace和R2D2。

V-trace負(fù)責(zé)從采樣的動(dòng)作中預(yù)測動(dòng)作的分布，R2D2負(fù)責(zé)根據(jù)動(dòng)作的預(yù)測未來值選擇一個(gè)動(dòng)作。

V-trace是基于策略梯度的方法，最早被IMPALA采用。由于Actor和Learner是異步執(zhí)行，而V-trace在異步體系架構(gòu)中的效果很好

第二種算法是R2D2，這是一種Q學(xué)習(xí)方法，DeepMind曾用這種算法將強(qiáng)化學(xué)習(xí)智能體在Atari游戲上的水平提高了4倍，并在52款游戲上超過了人類水平。

這種方法允許Q學(xué)習(xí)算法大規(guī)模硬件上運(yùn)行的同時(shí)仍然可以使用RNN。

實(shí)驗(yàn)結(jié)果

谷歌在DeepMind最近研究的開源足球游戲項(xiàng)目Google Research Football中進(jìn)行了基準(zhǔn)測試。

使用64個(gè)Cloud TPU核心實(shí)現(xiàn)了每秒240萬幀的數(shù)據(jù)傳輸速度，與之前的最新分布式IMPALA相比，提高了80倍。

IMPALA要想達(dá)到相同的速度，需要14000個(gè)CPU，而SEED RL只用了4160個(gè)CPU。對于相同的速度，IMPALA需要的CPU是SEED RL的3~4倍。

通過對并行計(jì)算的硬件加速器進(jìn)行優(yōu)化，我們就可以放心大膽地提高模型的大小。

在上面的足球游戲任務(wù)中，通過增加模型大小和輸入分辨率，可以解決從前未能解決的一些困難，讓訓(xùn)練模型的效率大幅提高。

傳送門

論文地址：

https://arxiv.org/abs/1910.06591

GitHub地址：

https://github.com/google-research/seed_rl