郭一璞 发自 凹非寺
量子位 报道 | 公众号 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#DeployingamodelwithTpuReplicator.
2repl=tf_replicator.TpuReplicator(
3num_workers=1,num_tpu_cores_per_worker=8
4)
5withrepl.context():
6model=resnet_model()
7base_optimizer=tf.train.AdamOptimizer()
8optimizer=repl.wrap_optimizer(base_optimizer)
9
10#...codetodefinereplicainput_fnandstep_fn.
11
12per_replica_loss=repl.run(step_fn,input_fn)
13train_op=tf.reduce_mean(per_replica_loss)
14
15withtf.train.MonitoredSession()assession:
16repl.init(session)
17foriinxrange(num_train_steps):
18session.run(train_op)
19repl.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文档
/alpha/guide/distribute_strategy
Colab笔记本
https://colab./github/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynb
GitHub笔记本
/tensorflow/docs/blob/master/site/en/r2/guide/distribute_strategy.ipynb
DeepMind博客
/blog/tf-replicator-distributed-machine-learning/
论文
/abs/1902.00465
— 完 —
订阅AI内参,获取AI行业资讯
购买AI书籍
诚挚招聘