怎麼把GPU上訓練的模型轉到TPU或者CPU上去？DeepMind發布新工具支招

雷鋒網 AI 科技評論按：DeepMind 有一支專門的科研平台團隊（the Research Platform Team），他們的職責是為 AI 學術研究構建加速計算的基礎設施。他們不經常亮相，但是這次由他們撰文介紹的 TF-Replicator 是一個極為有用的工具：它是又一個實用的軟體庫，可以幫助從未接觸過分散式系統的研究人員們輕鬆地在 GPU 集群和雲 TPU 集群上部署 TensorFlow 模型，也就成為了深度學習進行大規模工業化應用的重要組件。TF-Replicator 的程序化模型現在也已經作為 TensorFlow 的 tf.distribute.Strategy 的一部分開源在 https://www.tensorflow.org/alpha/guide/distribute_strategy。

DeepMind 科研平台團隊發表的技術博客中介紹了 TF-Replicator 的核心思想、開發經歷、技術挑戰。雷鋒網 AI 科技評論把主要內容編譯如下。

在 AlphaFold、BigGAN、AlphaStar 等近期的人工智慧科研成果中，我們很容易看到一個反覆出現的要素，那就是對輕鬆增加模型和計算能力規模的渴求。在越來越高的計算能力上訓練越來越大的模型，讓神經網站展現出了驚人的表現。為了應對這樣的需求，DeepMind 科研平台團隊開發了 TF-Replicator，它可以幫助科研人員輕鬆地為機器學習模型指定不同的硬體加速器、把運算拓展到更多設備上，以及在不同類型的加速器之間無縫切換。TF-Replicator 一開始是基於 TensorFlow 進行開發的，如今它的 API 也已經直接集成進了 TF 2.0 中。

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

DeepMind 科研平台團隊開發 TF-Replicator 的初始目標是為 DeepMind 的研究人員們提供一種簡單便捷的使用 TPU 的 API。不過 TensorFlow 原生的 TPU 計算 API 和 GPU 計算 API 有不少區別，這就為切換到 TPU 帶來了不小障礙。TF-Replicator 提供的 API 要更簡單、更友好，隱藏了 TF 原本的 TPU API 中的複雜之處。其中最關鍵的是，科研平台團隊和 DeepMind 之外的許多機器學習組織機構都進行了緊密的合作，以確保他們開發出的 TF-Replicator 有足夠的靈活性、易用性。

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

# 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)

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

對於多 GPU 配置中的運算，TF-Replicator 使用的是一種「圖內複製」（in-graph replication）的模式，每個設備的計算都會複製到同一個 TensorFlow 圖中。設備之間的溝通方式是連接不同設備的對應子圖的節點。在 TF-Replicator 中實現這個其實很有挑戰性，因為溝通可以發生在數據流圖中的任何一點上。運算的構建順序就變得尤其重要。

DeepMind 科研平台團隊最開始的想法是在單獨的 Python 線程中為每個設備同步構建各自的子圖。當遇到一個主動發起溝通的請求時，這些線程就會同步，主線程會插入執行所需的跨設備運算。在此之後，每個線程都會繼續構建它自己對應的設備的運算。然而，對這種方法進行仔細思考之後他們意識到 TF 中的圖構建 API 並不是線程安全的，這意味著同步在不同的線程中構建不同的子圖非常困難。所以他們轉而使用了圖重寫（graph rewriting），在所有設備的子圖都完成構建之後插入溝通。當這些子圖還在構建的時候，TF-Replicator 會在需要溝通的地方插入一個佔位符，子圖構建完成之後它就會在不同的設備間比對所有匹配的佔位符，然後把它們更換成正確的跨設備運算。

由於在設計和實現的全過程中都和學術界內的研究人員們保持緊密的合作，TF-Replicator 得以被構建成一個可以輕鬆地在多種硬體加速器之間拓展運算的庫，同時還讓用戶保留了前沿 AI 研究所需的高度控制和靈活性。

在開發完畢之前，TF-Replicator 就已經是 DeepMind 的研究人員們使用最多的 TPU 編程介面。TF-Replicator 的功能並不限制於訓練神經網路，不過它還是最常用來在大量數據上進行大批量訓練。比如 BigGAN 模型就是在最高達到 512 個 TPUv3 核心組成的集群上，以 2048 的批量大小進行訓練的。在帶有分散式的執行者-學習者設置的強化學習智能體中，更大規模的運算可以讓更多不同的執行者和環境進行交互，也可以產生新的經驗。TF-Replicator 已經為更多執行者的環境做好了準備，可以輕鬆地在多個硬體加速器上分布學習者（通常是神經網路）的運算。在論文中有更詳細的介紹。

DeepMind 科研平台團隊開發了許多有影響力的技術工具，TF-Replicator 只是其中之一。DeepMind 的許多深度學習科研成果，都有科研平台團隊的功勞。DeepMind 科研平台團隊在介紹分享自己的成果的同時，也歡迎對前沿 AI 研究感興趣的科研平台開源軟體工程師加入他們。

詳細介紹請見論文 https://arxiv.org/abs/1902.00465

via DeepMind Blog，雷鋒網 AI 科技評論編譯

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！