遷移單一工作站多 GPU 訓練

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

本指南示範如何將單一工作站多 GPU 工作流程從 TensorFlow 1 遷移至 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 環境定義中，而非在 tf.estimator.Estimator 中定義 config，才能例項化 Keras 模型、最佳化工具和指標。

如果您需要使用自訂訓練迴圈，請參閱「搭配自訂訓練迴圈使用 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 進行分散式訓練，請參閱下列文件