TensorFlow 1.11.0发布，一键多GPU（训练、预测和评价tf.keras模型）

导读：TensorFlow 1.11.0发布了，带来许多不错的新特性。其中最吸引人的特性应该算是tf.keras目前可以使用DistributionStrategy一键实现多GPU操作了。本文介绍在TensorFlow 1.11.0中如何用tf.keras和DistributionStrategy来一键实现多GPU训练、预测和评价。

tf.keras

---

tf.keras是TensorFlow早就有的特性，是keras在TensorFlow中集成的结果。由于其非常友好的接口、较低的入门门槛等因素，tf.keras成为了目前TensorFlow社区主推的特性之一。例如用下面简单的代码就可以实现一个网络的数据处理（如自动分batch）、网络构建、训练、评价等功能：

import tensorflow as tf
mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)

DistributionStrategy

---

DistributionStrategy API是构建多设备/机器训练的简单方式，开发者只需要在现有模型上做少量的修改，就可以用它们进行分布式训练。另外，DistributionStrategy在设计时考虑了同时兼容动态图(eager)和静态图。

目前TensorFlow支持三种DistributionStrategy：

• MirroredStrategy

• CollectiveAllReduceStrategy

• ParameterServerStrategy
  

本教程要使用的是MirroredStrategy，参数被同步到多个GPU中，在更新参数前，多个GPU对参数的更新操作被合并，因此在更新后，多个GPU中的参数依然能保持同步。

老版本中需要将kt.keras转换为Estimator才能使用DistributionStrategy

---

在老版本中，tf.keras不能直接使用DistributionStrategy，需要将其用下面的代码转换为Estimator，在Estimator中使用DistributionStrategy。

keras_estimator = keras.estimator.model_to_estimator(
  keras_model=model,
  config=config,
  model_dir='/tmp/model_dir')

在tf.keras中直接使用DistributionStrategy

---

最新的TensorFlow Github中给出了在tf.keras中直接使用DistributionStrategy的例子。

用tf.keras构建一个单层网络：

inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)

目前，使用DistributionStrategy需要使用tf.data.Dataset来作为数据输入：

features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
train_dataset = tf.data.Dataset.zip((features, labels))

这里我们为模型指定使用MirroredStrategy进行多GPU训练，代码非常简单：

distribution = tf.contrib.distribute.MirroredStrategy()
model.compile(loss='mean_squared_error',
              optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2),
              distribute=distribution)

使用常规的训练、评价和预测方法会自动在多GPU上进行：

model.fit(train_dataset, epochs=5, steps_per_epoch=10)
model.evaluate(eval_dataset)
model.predict(predict_dataset)

将tf.keras模型迁移到多GPU上运行只需要上面这些代码，它会自动切分输入、在每个设备（GPU）上复制层和变量、合并和更新梯度。

另外，模型的保存是无缝的，模型的保存和断点不受DistributionStrategy的影响。

-END-

专 · 知

人工智能领域26个主题知识资料全集获取与加入专知人工智能服务群: 欢迎微信扫一扫加入专知人工智能知识星球群，获取专业知识教程视频资料和与专家交流咨询！

请PC登录www.zhuanzhi.ai或者点击阅读原文，注册登录专知，获取更多AI知识资料！

请加专知小助手微信（扫一扫如下二维码添加），加入专知主题群（请备注主题类型：AI、NLP、CV、 KG等）交流~

 AI 项目技术 & 商务合作：bd@zhuanzhi.ai, 或扫描上面二维码联系！

请关注专知公众号，获取人工智能的专业知识！

点击“阅读原文”，使用专知