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_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