雷鋒網(wǎng) AI 科技評(píng)論按：DeepMind 有一支專門的科研平臺(tái)團(tuán)隊(duì)（the Research Platform Team），他們的職責(zé)是為 AI 學(xué)術(shù)研究構(gòu)建加速計(jì)算的基礎(chǔ)設(shè)施。他們不經(jīng)常亮相，但是這次由他們撰文介紹的 TF-Replicator 是一個(gè)極為有用的工具：它是又一個(gè)實(shí)用的軟件庫，可以幫助從未接觸過分布式系統(tǒng)的研究人員們輕松地在 GPU 集群和云 TPU 集群上部署 TensorFlow 模型，也就成為了深度學(xué)習(xí)進(jìn)行大規(guī)模工業(yè)化應(yīng)用的重要組件。TF-Replicator 的程序化模型現(xiàn)在也已經(jīng)作為 TensorFlow 的 tf.distribute.Strategy 的一部分開源在 https://www.tensorflow.org/alpha/guide/distribute_strategy。

DeepMind 科研平臺(tái)團(tuán)隊(duì)發(fā)表的技術(shù)博客中介紹了 TF-Replicator 的核心思想、開發(fā)經(jīng)歷、技術(shù)挑戰(zhàn)。雷鋒網(wǎng) AI 科技評(píng)論把主要內(nèi)容編譯如下。

在 AlphaFold、BigGAN、AlphaStar 等近期的人工智能科研成果中，我們很容易看到一個(gè)反復(fù)出現(xiàn)的要素，那就是對(duì)輕松增加模型和計(jì)算能力規(guī)模的渴求。在越來越高的計(jì)算能力上訓(xùn)練越來越大的模型，讓神經(jīng)網(wǎng)站展現(xiàn)出了驚人的表現(xiàn)。為了應(yīng)對(duì)這樣的需求，DeepMind 科研平臺(tái)團(tuán)隊(duì)開發(fā)了 TF-Replicator，它可以幫助科研人員輕松地為機(jī)器學(xué)習(xí)模型指定不同的硬件加速器、把運(yùn)算拓展到更多設(shè)備上，以及在不同類型的加速器之間無縫切換。TF-Replicator 一開始是基于 TensorFlow 進(jìn)行開發(fā)的，如今它的 API 也已經(jīng)直接集成進(jìn)了 TF 2.0 中。

TensorFlow 固然對(duì) CPU、GPU、TPU 都提供了直接的支持，但是用戶想要切換模型運(yùn)行在哪個(gè)設(shè)備上卻是一件很麻煩的事情。一般來說，這個(gè)過程都需要針對(duì)目標(biāo)硬件重新編寫專門的代碼，這也就限制了科研想法在具體平臺(tái)上的發(fā)揮。目前也有一些別的基于 TensorFlow 的框架在嘗試應(yīng)對(duì)這樣的問題，比如 Estimators，不過這些方案基本上都是針對(duì)生產(chǎn)環(huán)境的，缺少學(xué)術(shù)科研的快速迭代所需要的表達(dá)能力和靈活性。

構(gòu)建一個(gè)分布式機(jī)器學(xué)習(xí)庫

DeepMind 科研平臺(tái)團(tuán)隊(duì)開發(fā) TF-Replicator 的初始目標(biāo)是為 DeepMind 的研究人員們提供一種簡單便捷的使用 TPU 的 API。不過 TensorFlow 原生的 TPU 計(jì)算 API 和 GPU 計(jì)算 API 有不少區(qū)別，這就為切換到 TPU 帶來了不小障礙。TF-Replicator 提供的 API 要更簡單、更友好，隱藏了 TF 原本的 TPU API 中的復(fù)雜之處。其中最關(guān)鍵的是，科研平臺(tái)團(tuán)隊(duì)和 DeepMind 之外的許多機(jī)器學(xué)習(xí)組織機(jī)構(gòu)都進(jìn)行了緊密的合作，以確保他們開發(fā)出的 TF-Replicator 有足夠的靈活性、易用性。

TF-Replicator API

使用了 TF-Replicator 的代碼看起來和單臺(tái)設(shè)備上運(yùn)行的 TF 代碼非常相似，給了用戶足夠的自由來定義他們自己的模型運(yùn)行循環(huán)。用戶只需要做兩次簡單的定義，一，定義指向某個(gè)數(shù)據(jù)集的輸入函數(shù)；二，定義一個(gè)指定了模型邏輯的步驟函數(shù)（比如梯度下降中的某一步迭代）。代碼示例如下：

# Deploying a model with TpuReplicator.

repl = tf_replicator.TpuReplicator(

    num_workers=1, num_tpu_cores_per_worker=8

)

with repl.context():

  model = resnet_model()

  base_optimizer = tf.train.AdamOptimizer()

  optimizer = repl.wrap_optimizer(base_optimizer)

# ... code to define replica input_fn and step_fn.

per_replica_loss = repl.run(step_fn, input_fn)

train_op = tf.reduce_mean(per_replica_loss)

with tf.train.MonitoredSession() as session:

  repl.init(session)

  for i in xrange(num_train_steps):

    session.run(train_op)

  repl.shutdown(session)

想要把運(yùn)算拓展到多個(gè)設(shè)備，需要讓這些設(shè)備之間能夠相互溝通。在訓(xùn)練機(jī)器學(xué)習(xí)模型的情境中，最常見的溝通形式就是為隨機(jī)梯度下降 SGD 之類的優(yōu)化算法累積計(jì)算梯度。所以 DeepMind 科研平臺(tái)團(tuán)隊(duì)在 TF-Replicator 中用一個(gè)便捷的方法提供了 TensorFlow 各種優(yōu)化器的打包，這樣在更新模型的參數(shù)之前，不同設(shè)備分別計(jì)算的梯度就可以先進(jìn)行累積。對(duì)于其它的更通用化的溝通，TF-Replicator 提供了類似 MPI 的主動(dòng)式溝通模式，比如 `all_reduce` 和 `broadcast`。這些功能的幫助下，BigGAN 模型中使用到的全局批量歸一化（global batch normalisation）也可以簡單地在操作中實(shí)現(xiàn)，而這也是 BigGAN 增加訓(xùn)練規(guī)模非常重要的一步（可以參見 https://arxiv.org/abs/1809.11096 論文第三節(jié)）。

數(shù)據(jù)從主機(jī) host 送入不同的 GPU 中并馬上開始計(jì)算。當(dāng) GPU 之間需要交換信息的時(shí)候，它們會(huì)先同步再發(fā)送數(shù)據(jù)。

TF-Replicator 的實(shí)現(xiàn)

對(duì)于多 GPU 配置中的運(yùn)算，TF-Replicator 使用的是一種“圖內(nèi)復(fù)制”（in-graph replication）的模式，每個(gè)設(shè)備的計(jì)算都會(huì)復(fù)制到同一個(gè) TensorFlow 圖中。設(shè)備之間的溝通方式是連接不同設(shè)備的對(duì)應(yīng)子圖的節(jié)點(diǎn)。在 TF-Replicator 中實(shí)現(xiàn)這個(gè)其實(shí)很有挑戰(zhàn)性，因?yàn)闇贤梢园l(fā)生在數(shù)據(jù)流圖中的任何一點(diǎn)上。運(yùn)算的構(gòu)建順序就變得尤其重要。

DeepMind 科研平臺(tái)團(tuán)隊(duì)最開始的想法是在單獨(dú)的 Python 線程中為每個(gè)設(shè)備同步構(gòu)建各自的子圖。當(dāng)遇到一個(gè)主動(dòng)發(fā)起溝通的請(qǐng)求時(shí)，這些線程就會(huì)同步，主線程會(huì)插入執(zhí)行所需的跨設(shè)備運(yùn)算。在此之后，每個(gè)線程都會(huì)繼續(xù)構(gòu)建它自己對(duì)應(yīng)的設(shè)備的運(yùn)算。然而，對(duì)這種方法進(jìn)行仔細(xì)思考之后他們意識(shí)到 TF 中的圖構(gòu)建 API 并不是線程安全的，這意味著同步在不同的線程中構(gòu)建不同的子圖非常困難。所以他們轉(zhuǎn)而使用了圖重寫（graph rewriting），在所有設(shè)備的子圖都完成構(gòu)建之后插入溝通。當(dāng)這些子圖還在構(gòu)建的時(shí)候，TF-Replicator 會(huì)在需要溝通的地方插入一個(gè)占位符，子圖構(gòu)建完成之后它就會(huì)在不同的設(shè)備間比對(duì)所有匹配的占位符，然后把它們更換成正確的跨設(shè)備運(yùn)算。

當(dāng) TF-Replicator 構(gòu)建一個(gè)圖內(nèi)復(fù)制計(jì)算時(shí)，它會(huì)首先為每個(gè)設(shè)備分別單獨(dú)構(gòu)建運(yùn)算，然后把用戶指定了需要使用跨設(shè)備運(yùn)算的地方用占位符代替。當(dāng)所有設(shè)備的子圖都構(gòu)建完畢后，TF-Replicator 就把占位符替換成實(shí)際的跨設(shè)備運(yùn)算，最終連接了不同的設(shè)備。

在 DeepMind 構(gòu)建一個(gè)基礎(chǔ)平臺(tái)

由于在設(shè)計(jì)和實(shí)現(xiàn)的全過程中都和學(xué)術(shù)界內(nèi)的研究人員們保持緊密的合作，TF-Replicator 得以被構(gòu)建成一個(gè)可以輕松地在多種硬件加速器之間拓展運(yùn)算的庫，同時(shí)還讓用戶保留了前沿 AI 研究所需的高度控制和靈活性。

在開發(fā)完畢之前，TF-Replicator 就已經(jīng)是 DeepMind 的研究人員們使用最多的 TPU 編程接口。TF-Replicator 的功能并不限制于訓(xùn)練神經(jīng)網(wǎng)絡(luò)，不過它還是最常用來在大量數(shù)據(jù)上進(jìn)行大批量訓(xùn)練。比如 BigGAN 模型就是在最高達(dá)到 512 個(gè) TPUv3 核心組成的集群上，以 2048 的批量大小進(jìn)行訓(xùn)練的。在帶有分布式的執(zhí)行者-學(xué)習(xí)者設(shè)置的強(qiáng)化學(xué)習(xí)智能體中，更大規(guī)模的運(yùn)算可以讓更多不同的執(zhí)行者和環(huán)境進(jìn)行交互，也可以產(chǎn)生新的經(jīng)驗(yàn)。TF-Replicator 已經(jīng)為更多執(zhí)行者的環(huán)境做好了準(zhǔn)備，可以輕松地在多個(gè)硬件加速器上分布學(xué)習(xí)者（通常是神經(jīng)網(wǎng)絡(luò)）的運(yùn)算。在論文中有更詳細(xì)的介紹。

DeepMind 科研平臺(tái)團(tuán)隊(duì)開發(fā)了許多有影響力的技術(shù)工具，TF-Replicator 只是其中之一。DeepMind 的許多深度學(xué)習(xí)科研成果，都有科研平臺(tái)團(tuán)隊(duì)的功勞。DeepMind 科研平臺(tái)團(tuán)隊(duì)在介紹分享自己的成果的同時(shí)，也歡迎對(duì)前沿 AI 研究感興趣的科研平臺(tái)開源軟件工程師加入他們。

詳細(xì)介紹請(qǐng)見論文 https://arxiv.org/abs/1902.00465

via DeepMind Blog，雷鋒網(wǎng) AI 科技評(píng)論編譯

雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。