迁移单工作进程多 GPU 训练

##### Copyright 2021 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

迁移单工作进程多 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 进行单工作进程分布式训练#

此示例演示了单工作进程多 GPU 训练的 TensorFlow 1 规范工作流。您需要通过 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 进行分布式训练，请查看以下文档：