迁移单工作器多 GPU 训练

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

本指南演示了如何将 TensorFlow 1 中的单工作器多 GPU 工作流程迁移到 TensorFlow 2。

在单台机器上的多个 GPU 上执行同步训练

设置

从导入和简单的演示数据集开始

import tensorflow as tf
import tensorflow.compat.v1 as tf1

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

TensorFlow 1：使用 tf.estimator.Estimator 进行单工作器分布式训练

此示例演示了 TensorFlow 1 中单工作器多 GPU 训练的规范工作流程。您需要通过 tf.estimator.Estimator 的 config 参数设置分布式策略 (tf.distribute.MirroredStrategy)

def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

strategy = tf1.distribute.MirroredStrategy()
config = tf1.estimator.RunConfig(
    train_distribute=strategy, eval_distribute=strategy)
estimator = tf1.estimator.Estimator(model_fn=_model_fn, config=config)

train_spec = tf1.estimator.TrainSpec(input_fn=_input_fn)
eval_spec = tf1.estimator.EvalSpec(input_fn=_eval_input_fn)
tf1.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

TensorFlow 2：使用 Keras 进行单工作器训练

迁移到 TensorFlow 2 时，您可以使用 Keras API 与 tf.distribute.MirroredStrategy。

如果您使用 tf.keras API 进行模型构建，并使用 Keras Model.fit 进行训练，主要区别在于在 Strategy.scope 的上下文中实例化 Keras 模型、优化器和指标，而不是为 tf.estimator.Estimator 定义 config。

如果您需要使用自定义训练循环，请查看 使用 tf.distribute.Strategy 与自定义训练循环 指南。

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss='mse')
model.fit(dataset)
model.evaluate(eval_dataset, return_dict=True)

后续步骤

要详细了解 TensorFlow 2 中使用 tf.distribute.MirroredStrategy 进行分布式训练，请查看以下文档