在 Cloud TPU 上训练 ResNet (TF 2.x)

本教程介绍如何使用 tf.distribute.TPUStrategy 在 Cloud TPU 上训练 Keras ResNet 模型。

如果您不熟悉 Cloud TPU，强烈建议您浏览快速入门，了解如何创建 TPU 和 Compute Engine 虚拟机。

警告：本教程使用第三方数据集。Google 不对此数据集的有效性或任何其他方面提供任何代理、担保或其他保证。

目标

• 创建 Cloud Storage 存储分区以保存数据集和模型输出。
• 准备与 ImageNet 数据集类似的 fake_imagenet 数据集。
• 运行训练作业。
• 验证输出结果。

费用

本教程使用 Google Cloud 的以下收费组件：

• Compute Engine
• Cloud TPU
• Cloud Storage

请使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用。

准备工作

在开始学习本教程之前，请检查您的 Google Cloud 项目是否已正确设置。

1. 登录您的 Google Cloud 帐号。如果您是 Google Cloud 新手，请创建一个帐号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金，用于运行、测试和部署工作负载。

2. 在 Google Cloud Console 的项目选择器页面上，选择或创建一个 Google Cloud 项目。

   注意：如果您不打算保留在此过程中创建的资源，请创建新的项目，而不要选择现有的项目。完成本教程介绍的步骤后，您可以删除所创建的项目，并移除与该项目关联的所有资源。

   转到“项目选择器”

3. 确保您的 Cloud 项目已启用结算功能。 了解如何确认您的项目是否已启用结算功能。

本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源，以免产生不必要的费用。

设置资源

本部分介绍如何为教程设置 Cloud Storage 存储空间、虚拟机和 Cloud TPU 资源。

重要提示 ：在同一地区/区域中设置所有资源（Compute Engine 虚拟机、Cloud TPU 和 Cloud Storage 存储分区），可降低网络延迟和网络费用。虚拟机和 TPU 节点位于特定地区，即区域内的细分。

1. 打开一个 Cloud Shell 窗口。

   打开 Cloud Shell

2. 为项目 ID 创建一个变量。

   export PROJECT_ID=project-id
   

3. 配置 gcloud 命令行工具，以使用要在其中创建 Cloud TPU 的项目。

   gcloud config set project ${PROJECT_ID}
   

   当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 GCP API 调用。

4. 为 Cloud TPU 项目创建服务帐号。

   gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
   

   该命令将返回以下格式的 Cloud TPU 服务帐号：

   [e-mail]
   

5. 使用以下命令创建 Cloud Storage 存储分区：

   gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
   

   此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的 gcloud compute tpus execution-groups 工具会为您在上一步中设置的 Cloud TPU 服务帐号设置默认权限。如果您需要更精细的权限，请查看访问级层权限。

   存储分区位置必须要与 Compute Engine（虚拟机）和 Cloud TPU 节点位于同一区域。

6. 使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。

   $ gcloud compute tpus execution-groups create \
    --vm-only \
    --name=resnet-tutorial \
    --zone=europe-west4-a \
    --disk-size=300 \
    --machine-type=n1-standard-16 \
    --tf-version=2.5.0
   

   命令标志说明

   vm-only

   仅创建虚拟机。默认情况下，gcloud compute tpus execution-groups 命令会同时创建虚拟机和 Cloud TPU。

   name

   要创建的 Cloud TPU 的名称。

   zone

   拟在其中创建 Cloud TPU 的区域。

   disk-size

   由 gcloud compute tpus execution-groups 命令创建的虚拟机的硬盘大小（以 GB 为单位）。

   machine-type

   要创建的 Compute Engine 虚拟机的机器类型。

   tf-version

   虚拟机上的 TensorFlow gcloud compute tpus execution-groups 安装的版本。

   如需详细了解 gcloud 命令，请参阅 gcloud 参考文档。

7. 在出现提示时，按 y 创建 Cloud TPU 资源。

   注意：第一次在项目上运行 gcloud compute tpus execution-groups 时，大约需要 5 分钟来执行一些启动任务，例如 SSH 密钥传播和 API 启动。

   当 gcloud compute tpus execution-groups 命令执行完毕后，验证 shell 提示符已从 username@projectname 更改为 username@vm-name。 此变化表明您现已登录 Compute Engine 虚拟机。

   注意：如果未连接到 Compute Engine 实例，您可以通过运行以下命令进行连接：

   gcloud compute ssh resnet-tutorial --zone=europe-west4-a
   

   继续按照这些说明操作，在 Compute Engine 实例中运行以 (vm)$ 开头的每个命令。

设置 Cloud Storage 存储分区变量

设置以下环境变量，将 bucket-name 替换为 Cloud Storage 存储分区的名称：

(vm)$ export STORAGE_BUCKET=gs://bucket-name

(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
(vm)$ export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models/"

训练应用应该能够访问 Cloud Storage 中的训练数据。在训练期间，训练应用还会使用您的 Cloud Storage 存储分区来存储检查点。

使用 fake_imagenet 来训练和评估 ResNet 模型

ImageNet 是一个图像数据库。数据库中的图片被整理为一个层次结构，该层次结构中的每个节点由成百上千个图片进行描述。

本教程使用演示版的完整 ImageNet 数据集，该数据集又称为 fake_imagenet。此演示版本可用于测试教程，同时降低通常与使用完整 ImageNet 数据库运行模型相关的存储和时间要求。

此 fake_imagenet 数据集位于 Cloud Storage 上的以下位置：

gs://cloud-tpu-test-datasets/fake_imagenet

此 fake_imagenet 数据集仅用于了解如何使用 Cloud TPU 并验证端到端性能。准确率数字和保存的模型并无实际意义。

如需了解如何下载和处理完整 ImageNet 数据集，请参阅下载、预处理和上传 ImageNet 数据集。

注意：对于本教程，请确保不要将 STORAGE_BUCKET 环境变量设置为 fake_imagenet 数据集的路径。您可以从 gs://cloud-tpu-test-datasets 读取，但无法对其执行写入操作。因此，您无法用其写出训练日志。确保将 STORAGE_BUCKET 环境变量设置为您自己的 Cloud Storage 存储分区，如上所示。

注意：如果想要监控模型的输出和性能，请参照设置 TensorBoard 指南。

1. 使用 gcloud 命令启动 Cloud TPU 资源。

   (vm)$ gcloud compute tpus execution-groups create \
    --tpu-only \
    --accelerator-type=v3-8  \
    --name=resnet-tutorial \
    --zone=europe-west4-a \
    --tf-version=2.5.0
   

   命令标志说明

   tpu-only

   创建 Cloud TPU，而不创建虚拟机。默认情况下，gcloud compute tpus execution-groups 命令会同时创建虚拟机和 Cloud TPU。

   accelerator-type

   要创建的 Cloud TPU 的类型。

   name

   要创建的 Cloud TPU 的名称。

   zone

   拟在其中创建 Cloud TPU 的区域。

   tf-version

   要在虚拟机上安装的 TensorFlow 的版本。

2. 设置 Cloud TPU 名称变量。

   (vm)$ export TPU_NAME=resnet-tutorial
   

3. ResNet 训练脚本需要额外的软件包。立即安装：

   (vm)$ sudo pip3 install tensorflow-model-optimization>=0.1.3
   

4. 导航到 ResNet-50 模型目录：

   (vm)$ cd /usr/share/models/official/vision/image_classification/resnet/
   

   注意：如果您运行的是 TensorFlow 2.1，请使用路径：/usr/share/models/official/vision/image_classification/

5. 运行训练脚本。使用一个 false_imagenet 数据集并用一个周期训练 ResNet。

   (vm)$ python3 resnet_ctl_imagenet_main.py \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --data_dir=${DATA_DIR} \
    --batch_size=1024 \
    --steps_per_loop=500 \
    --train_epochs=1 \
    --use_synthetic_data=false \
    --dtype=fp32 \
    --enable_eager=true \
    --enable_tensorboard=true \
    --distribution_strategy=tpu \
    --log_steps=50 \
    --single_l2_loss_op=true \
    --use_tf_function=true
   

   命令标志说明

   tpu

   Cloud TPU 的名称。

   model_dir

   在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。

   data_dir

   训练输入的 Cloud Storage 路径。在此示例中，该路径设置为 fake_imagenet 数据集。

   batch_size

   训练批次大小。

   steps_per_loop

   将状态保存到 CPU 之前要运行的训练步数。 一个训练步骤是处理一批次的样本。这包括前向传导和反向传播。

   train_epochs

   训练模型的周期数。

   use_synthetic_data

   设置为 true 以使用合成数据进行训练。

   dtype

   用于训练的数据类型。

   enable_eager

   启用 TensorFlow 即刻执行。

   enable_tensorboard

   启用 TensorBoard。

   distribution_strategy

   设置为 tpu 以在 Cloud TPU 上训练 ResNet 模型。

   log_steps

   记录时间信息（例如 examples per second）之前要执行的训练步数。

   single_l2_loss_op

   设置为 true 以计算关联权重的 L2_loss，而不是使用 Keras 每层 L2 损失。

   use_tf_function

   设置为 true 以将训练和测试步骤封装在 tf.function 内。

ResNet 1 周期的训练将在 10 分钟内在 v3-8 TPU 节点上完成。训练脚本输出应包括如下所示的文本：

{
  'train_loss': 1.435225,
  'train_accuracy': 0.00084427913
}

训练脚本也会执行评估。评估输出应包含如下文本：

Run stats:
{
  'eval_loss': 0.861013,
  'eval_acc': 0.001,
  'train_loss': 1.435225,
  'train_acc': 0.00084427913,
  'step_timestamp_log': [
    'BatchTimestamp<batch_index: 0,
    timestamp: 1606330585.7613473>',
    'BatchTimestamp<batch_index: 500,
    timestamp: 1606330883.8486104>',
    'BatchTimestamp<batch_index: 1000,
    timestamp: 1606331119.515312>',
    'BatchTimestamp<batch_index: 1251,
    timestamp: 1606331240.7516596>'
  ],
  'train_finish_time': 1606331296.395158,
  'avg_exp_per_second': 1951.6983246161021
}

要训练 ResNet 收敛，请运行 90 个周期，如以下脚本所示。训练和评估是一起完成的。每个周期都有 1251 个步骤，总计 112590 个训练步骤和 48 个评估步骤。

(vm)$ python3 resnet_ctl_imagenet_main.py \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --data_dir=${DATA_DIR} \
    --batch_size=1024 \
    --steps_per_loop=500 \
    --train_epochs=90 \
    --use_synthetic_data=false \
    --dtype=fp32 \
    --enable_eager=true \
    --enable_tensorboard=true \
    --distribution_strategy=tpu \
    --log_steps=50 \
    --single_l2_loss_op=true \
    --use_tf_function=true

命令标志说明

tpu

Cloud TPU 的名称。

model_dir

在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。

data_dir

训练输入的 Cloud Storage 路径。在此示例中，该路径设置为 fake_imagenet 数据集。

batch_size

训练批次大小。

steps_per_loop

将状态保存到 CPU 之前要运行的训练步数。训练步骤是对一批样本的处理。包括前向传导和反向传播。

train_epochs

训练模型的周期数。

use_synthetic_data

是否使用合成数据进行训练。

dtype

用于训练的数据类型。

enable_eager

启用 TensorFlow 即刻执行。

enable_tensorboard

启用 TensorBoard。

distribution_strategy

设置为 tpu 以在 Cloud TPU 上训练 ResNet 模型。

log_steps

记录时间信息（例如 examples per second）之前要执行的训练步数。

single_l2_loss_op

设置为 true 以计算关联权重的 L2_loss，而不是使用 Keras 每层 L2 损失。

use_tf_function

设置为 true 以将训练和测试步骤封装在 tf.function 内。

由于训练和评估是利用 fake_imagenet 数据集完成的，因此输出结果不会反映实际输出，实际输出在利用实际数据集执行训练和评估时才会出现。

此时，您可以结束本教程并清理 GCP 资源，也可以进一步了解如何在 Cloud TPU Pod 上运行模型。

使用 Cloud TPU Pod 扩缩模型

您可以使用 Cloud TPU Pod 扩缩模型，以便更快获得结果。完全受支持的 ResNet-50 模型可与以下 Pod 切片配合使用：

• v2-32
• v3-32

借助 Cloud TPU Pod，训练和评估可同时完成。

使用 Cloud TPU Pod 进行训练

1. 删除为在单台设备上训练模型而创建的 Cloud TPU 资源。

   (vm)$ gcloud compute tpus execution-groups delete resnet-tutorial \
     --zone=europe-west4-a \
     --tpu-only

2. 运行 gcloud compute tpus execution-groups 命令，并使用 accelerator-type 参数指定要使用的 Pod 切片。例如，以下命令使用 v3-32 Pod 切片。

   (vm)$ gcloud compute tpus execution-groups  create --name=resnet-tutorial \
     --accelerator-type=v3-32  \
     --zone=europe-west4-a \
     --tf-version=2.5.0 \
     --tpu-only
   

   命令标志说明

   name

   要创建的 Cloud TPU 的名称。

   accelerator-type

   要创建的 Cloud TPU 的类型。

   zone

   拟在其中创建 Cloud TPU 的区域。

   tf-version

   在虚拟机上安装的 Tensorflow gcloud 的版本。

   tpu-only

   仅创建 Cloud TPU。默认情况下，gcloud 命令会同时创建虚拟机和 Cloud TPU。

3. 设置一些所需的环境变量：

   (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x-pod
   

4. 训练模型。

   (vm)$ python3 resnet_ctl_imagenet_main.py \
     --tpu=${TPU_NAME} \
     --model_dir=${MODEL_DIR} \
     --data_dir=${DATA_DIR} \
     --batch_size=4096 \
     --steps_per_loop=500 \
     --train_epochs=1 \
     --use_synthetic_data=false \
     --dtype=fp32 \
     --enable_eager=true \
     --enable_tensorboard=true \
     --distribution_strategy=tpu \
     --log_steps=50 \
     --single_l2_loss_op=true \
     --use_tf_function=true

   命令标志说明

   tpu

   Cloud TPU 的名称。

   model_dir

   在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。

   data_dir

   训练输入的 Cloud Storage 路径。在此示例中，该路径设置为 fake_imagenet 数据集。

   batch_size

   训练批次大小。

   steps_per_loop

   将状态保存到 CPU 之前要运行的训练步数。 一个训练步骤是处理一批次的样本。包括前向传导和反向传播。

   train_epochs

   训练模型的周期数。

   use_synthetic_data

   设置为 true 以使用合成数据进行训练。

   dtype

   用于训练的数据类型。

   enable_eager

   设置为 true 以启用 TensorFlow 即刻执行。

   enable_tensorboard

   设置为 true 以启用 TensorBoard。

   distribution_strategy

   设置为 tpu 以在 TPU 上训练 ResNet 模型。

   log_steps

   记录时间信息（例如 examples per second）之前要执行的训练步数。

   single_l2_loss_op

   计算串联权重的 L2_loss，而不是使用 Keras 每层 L2 损失。

   use_tf_function

   将训练和测试步骤封装在 tf.function 内。

该程序将基于 fake_imagenet 数据集对模型进行 1 个周期的训练（共 312 个训练步骤和 12 个评估步骤）。此训练在 v3-32 Cloud TPU 上大约需要 2 分钟。训练脚本应显示如下所示的文本：

step: 312
steps_per_second: 1.38
{
  'train_loss': 0.3755003,
  'train_accuracy': 0.000983605
}

该脚本还会执行评估，并且应显示如下文本：

step: 312
evaluation metric: {
  'test_loss': 0.20280758,
  'test_accuracy': 0.001
}

清理

为避免因本教程中使用的资源导致您的 Google Cloud 帐号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

1. 断开与 Compute Engine 实例的连接（如果您尚未这样做）：

   (vm)$ exit
   

   您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

2. 在您的 Cloud Shell 中，使用以下命令删除您的 Compute Engine 虚拟机和 Cloud TPU：

   $ gcloud compute tpus execution-groups delete resnet-tutorial \
     --zone=europe-west4-a
   

3. 通过运行 gcloud compute tpus execution-groups list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

   $ gcloud compute tpus execution-groups list \
    --zone=europe-west4-a
   

   您应该会看到如下所示的空白 TPU 列表：

      NAME             STATUS
   

4. 使用 gsutil 删除 Cloud Storage 存储分区，如下所示。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称。

   $ gsutil rm -r gs://bucket-name
   

注意：如需了解免费存储限制和其他价格信息，请参阅 Cloud Storage 价格指南。

后续步骤

在本教程中，您已使用示例数据集训练 RESNET 模型。此训练的结果（在大多数情况下）不能用于推断。要使用模型进行推断，您可以在公开提供的数据集或您自己的数据集上训练数据。在 Cloud TPU 上训练的模型要求数据集采用 TFRecord 格式。

您可以使用数据集转换工具示例将图片分类数据集转换为 TFRecord 格式。如果您未使用图片分类模型，则必须自行将数据集转换为 TFRecord 格式。如需了解详情，请参阅 TFRecord 和 tf.Example

超参数调节

如需使用数据集提升模型的性能，您可以调节模型的超参数。您可以在 GitHub 上寻找所有 TPU 支持模型通用的超参数的相关信息。您可以在每个模型的源代码中寻找模型专用超参数的相关信息。如需详细了解超参数调节，请参阅超参数调节概览、使用超参数调节服务和调节超参数。

推理

训练模型后，您可以使用该模型进行推断（也称为预测）。AI Platform 是一款基于云的解决方案，用于开发、训练和部署机器学习模型。部署模型后，您可以使用 AI Platform Prediction 服务。

• 遵循数据集转换教程，了解如何使用自己的数据替代 fake_imagenet 或 ImageNet 数据集来进行训练和评估。该教程介绍了如何使用图片分类数据转换器示例脚本，将用于图片分类的原始数据集转换为 Cloud TPU Tensorflow 模型可用的 TFRecord。
• 运行 Cloud TPU colab，演示如何使用您自己的图片数据运行图片分类模型。
• 浏览其他 Cloud TPU 教程。
• 学习 TensorBoard 中 TPU 监控工具的使用方法。
• 了解如何使用 Cloud TPU 和 GKE 训练 ResNet。