从 TPUEstimator 迁移到 TPUStrategy

##### 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.

从 TPUEstimator 迁移到 TPUStrategy#

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

本指南演示了如何将在 TPU 上运行的工作流从 TensorFlow 1 的 TPUEstimator API 迁移到 TensorFlow 2 的 TPUStrategy API。

  • 在 TensorFlow 1 中，可以使用 tf.compat.v1.estimator.tpu.TPUEstimator API 训练和评估模型，以及在 (Cloud) TPU 上执行推断和保存模型（用于应用）。

  • 在 TensorFlow 2 中，要在 TPU 和 TPU Pod（通过专用高速网络接口连接的 TPU 设备集合）上进行同步训练，需要使用 TPU 分布策略 tf.distribute.TPUStrategy。此策略适用于 Keras API（包括用于模型构建 (tf.keras.Model)、优化器 (tf.keras.optimizers.Optimizer) 和训练 (Model.fit) 的 API）以及自定义训练循环（使用 tf.function 和 tf.GradientTape）。

对于端到端 TensorFlow 2 示例，请查看使用 TPU 指南（即 TPU 上的分类部分）和在 TPU 上使用 BERT 解决 GLUE 任务教程。您可能还会发现分布式训练指南十分有用，它涵盖了包括 TPUStrategy 在内的所有 TensorFlow 分布策略。

安装#

从导入和用于演示目的的简单数据集开始：

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

features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]

TensorFlow 1：使用 TPUEstimator 在 TPU 上驱动模型#

指南的这一部分演示了如何在 TensorFlow 1 中使用 tf.compat.v1.estimator.tpu.TPUEstimator 执行训练和评估。

要使用 TPUEstimator，首先定义几个函数：训练数据的输入函数，评估数据的评估输入函数，以及告知 TPUEstimator 如何使用特征和标签定义训练运算的模型函数：

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((features, labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((eval_features, eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  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.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

定义这些函数后，创建一个提供聚簇信息的 tf.distribute.cluster_resolver.TPUClusterResolver 和一个 tf.compat.v1.estimator.tpu.RunConfig 对象。连同您定义的模型函数，现在可以创建一个 TPUEstimator。在这里，您将通过跳过检查点保存来简化流程。随后，您将为 TPUEstimator 的训练和评估指定批次大小。

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))

tpu_config = tf1.estimator.tpu.TPUConfig(iterations_per_loop=10)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn,
    config=config,
    train_batch_size=8,
    eval_batch_size=8)

调用 TPUEstimator.train 以开始训练模型：

estimator.train(_input_fn, steps=1)

然后，调用 TPUEstimator.evaluate 以使用评估数据评估模型：

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2：使用 Keras Model.fit 和 TPUStrategy 在 TPU 上驱动模型#

在 TensorFlow 2 中，要在 TPU 工作进程上进行训练，请将 tf.distribute.TPUStrategy 与 Keras API 一起用于模型定义和训练/评估。（有关使用 Keras Model.fit 和自定义训练循环（使用 tf.function 和 tf.GradientTape）进行训练的更多示例，请参阅使用 TPU 指南。）

由于您需要执行一些初始化工作以连接到远程聚簇并初始化 TPU 工作进程，因此首先创建一个 TPUClusterResolver 以提供聚簇信息并连接到聚簇。（在使用 TPU 指南的 TPU 初始化部分中了解详情。）

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

接下来，准备好数据后，将创建一个 TPUStrategy，然后在此策略的范围下定义模型、指标和优化器。

要实现与 TPUStrategy 相当的训练速度，应当确保在 Model.compile 中为 steps_per_execution 选择一个数字，因为它指定了每次 tf.function 调用期间要运行的批次数，并且对性能至关重要。此参数类似于 TPUEstimator 中使用的 iterations_per_loop。如果您使用自定义训练循环，则应确保在 tf.function 修饰过的训练函数中运行多个步骤。如需了解详情，请转到使用 TPU 指南的在 tf.function 中通过多个步骤提升性能部分。

tf.distribute.TPUStrategy 可以支持有界动态形状，即可以推断出动态形状计算的上限。但是，与静态形状相比，动态形状可能会引入一些性能开销。因此，如果有可能，通常建议将输入形状设为静态，尤其是在训练中。返回动态形状的一个常见运算是 tf.data.Dataset.batch(batch_size)，因为流中剩余的样本数可能小于批次大小。因此，在 TPU 上进行训练时，应使用 tf.data.Dataset.batch(..., drop_remainder=True) 来获得最佳训练性能。

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).shuffle(10).repeat().batch(
        8, drop_remainder=True).prefetch(2)
eval_dataset = tf.data.Dataset.from_tensor_slices(
    (eval_features, eval_labels)).batch(1, drop_remainder=True)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
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, "mse", steps_per_execution=10)

这样，您就可以使用训练数据集训练模型了：

model.fit(dataset, epochs=5, steps_per_epoch=10)

最后，使用评估数据集评估模型：

model.evaluate(eval_dataset, return_dict=True)

后续步骤#

要详细了解 TensorFlow 2 中的 TPUStrategy，请查看以下资源：

要详细了解自定义训练，请参阅：

TPU（Google 用于机器学习的专用 ASIC）可通过 Google Colab、TPU Research Cloud 和 Cloud TPU 获得。