DeepMind最近為TensorFlow 2.0獻祭了自己私藏的工具：
TF-Replicator，本來是內(nèi)部自用的一個軟件庫，能夠讓從來沒做過分布式系統(tǒng)的研究人員方便地在多GPU/云TPU上部署他們的TensorFlow模型，也適用于Keras。
目前，TF-Replicator的編程模型已經(jīng)作為TensorFlow中tf.distribute.Strategy的一部分開源。
推特上的一位工程師驚嘆：這簡直是TensorFlow 2.0里隱藏的寶藏啊！
使用TF-Replicator編寫的代碼與TensorFlow中為單個設(shè)備編寫的代碼類似，允許用戶自由定義自己的模型運行循環(huán)。
用戶只需要定義兩個部分：
1.公開數(shù)據(jù)集的輸入函數(shù)；
2.模型邏輯的步驟函數(shù)。
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)
現(xiàn)在，我們用GAN來測試一下TF-Replicator的效果。這里用到的是在ImageNet上訓(xùn)練的譜歸一化GAN（SN-GAN, arXiv:1802.05957）。
相比在單一的一塊GPU上訓(xùn)練，用TF-Replicator在多塊GPU上分布式訓(xùn)練的效果要好得多。
比如，生成橙子的圖片，這是batch size 8和batch size 16的時候：
基本看不出來是橙子了。
batch size 32和batch size 64要好一些，能看出來是橙子，但是一個像長了毛，一個像被拍了一巴掌：
batch size 128有了橙子果肉，batch size 256形狀相對正常了：
示例中最高的batch size 512，橙子的形狀已經(jīng)和真實的橙子差不多了，果肉和果肉瓣之間的白色也可以看出來，除了皮有點厚之外這橙子質(zhì)量沒問題。
從分數(shù)來看，只要將batch size從64提高到512就可以將出實得分提高大約50%。
效果不錯，希望DeepMind繼續(xù)公開一些自用好貨。
最后，附上官方的相關(guān)文檔：
TensorFlow文檔
https://www.tensorflow.org/alpha/guide/distribute_strategy
Colab筆記本
https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynb
GitHub筆記本
https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynb
DeepMind博客
https://deepmind.com/blog/tf-replicator-distributed-machine-learning/
論文
https://arxiv.org/abs/1902.00465
— 完 —
訂閱AI內(nèi)參，獲取AI行業(yè)資訊
購買AI書籍
誠摯招聘
