DeepMind私貨公開，推出分散式機器學習庫，TensorFlow、Keras可用
郭一璞 發自 凹非寺
量子位 報道 | 公眾號 QbitAI
DeepMind最近為TensorFlow 2.0獻祭了自己私藏的工具：
TF-Replicator，本來是內部自用的一個軟體庫，能夠讓從來沒做過分散式系統的研究人員方便地在多GPU/雲TPU上部署他們的TensorFlow模型，也適用於Keras。
目前，TF-Replicator的編程模型已經作為TensorFlow中tf.distribute.Strategy的一部分開源。
推特上的一位工程師驚嘆：這簡直是TensorFlow 2.0里隱藏的寶藏啊！
怎麼用
使用TF-Replicator編寫的代碼與TensorFlow中為單個設備編寫的代碼類似，允許用戶自由定義自己的模型運行循環。
用戶只需要定義兩個部分：
1.公開數據集的輸入函數；
2.模型邏輯的步驟函數。
1# Deploying a model with TpuReplicator.
2repl = tf_replicator.TpuReplicator(
3 num_workers=1, num_tpu_cores_per_worker=8
4)
5with repl.context():
6 model = resnet_model()
7 base_optimizer = tf.train.AdamOptimizer()
8 optimizer = repl.wrap_optimizer(base_optimizer)
9
10# ... code to define replica input_fn and step_fn.
11
12per_replica_loss = repl.run(step_fn, input_fn)
13train_op = tf.reduce_mean(per_replica_loss)
14
15with tf.train.MonitoredSession() as session:
16 repl.init(session)
17 for i in xrange(num_train_steps):
18 session.run(train_op)
19 repl.shutdown(session)
拿來GAN一下試試
現在，我們用GAN來測試一下TF-Replicator的效果。這裡用到的是在ImageNet上訓練的譜歸一化GAN（SN-GAN, arXiv:1802.05957）。
相比在單一的一塊GPU上訓練，用TF-Replicator在多塊GPU上分散式訓練的效果要好得多。
比如，生成橙子的圖片，這是batch size 8和batch size 16的時候：
基本看不出來是橙子了。
batch size 32和batch size 64要好一些，能看出來是橙子，但是一個像長了毛，一個像被拍了一巴掌：
batch size 128有了橙子果肉，batch size 256形狀相對正常了：
示例中最高的batch size 512，橙子的形狀已經和真實的橙子差不多了，果肉和果肉瓣之間的白色也可以看出來，除了皮有點厚之外這橙子質量沒問題。
從分數來看，只要將batch size從64提高到512就可以將出實得分提高大約50%。
效果不錯，希望DeepMind繼續公開一些自用好貨。
傳送門
最後，附上官方的相關文檔：
TensorFlow文檔
https://www.tensorflow.org/alpha/guide/distribute_strategyColab筆記本
https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynbGitHub筆記本
https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynbDeepMind博客
https://deepmind.com/blog/tf-replicator-distributed-machine-learning/論文
https://arxiv.org/abs/1902.00465— 完 —
量子位 · QbitAI
?? ? 追蹤AI技術和產品新動態
戳右上角「+關注」獲取最新資訊↗↗
如果喜歡，請分享or點贊吧~比心?
推薦閱讀：
