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

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

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

目标

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

费用

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

Compute Engine
Cloud TPU
Cloud Storage

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

准备工作

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

登录您的 Google 帐号。
如果您还没有 Google 帐号，请注册一个新帐号。

在 Cloud Console 的项目选择器页面上，选择或创建 Cloud 项目。

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

转到项目选择器页面

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

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

设置资源

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

打开一个 Cloud Shell 窗口。

打开 Cloud Shell
为项目 ID 创建一个变量。
```
export PROJECT_ID=project-id
```
配置 gcloud 命令行工具，以使用要在其中创建 Cloud TPU 的项目。
```
gcloud config set project ${PROJECT_ID}
```

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

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

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

service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com

使用以下命令创建 Cloud Storage 存储分区：
```
gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
```
此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的 ctpu up 工具会为您在上一步中设置的 Cloud TPU 服务帐号设置默认权限。如果您需要更精细的权限，请查看访问级层权限。

存储分区位置必须要与 Compute Engine（虚拟机）和 Cloud TPU 节点位于同一地区。
使用 ctpu up 命令启动 Compute Engine 虚拟机资源。
```
ctpu up --zone=europe-west4-a \
 --vm-only \
 --name=resnet-tutorial \
 --disk-size-gb=300 \
 --machine-type=n1-standard-16 \
 --tf-version=2.3.1
```
注意：如果您有多个项目，则必须使用 --project 标志指定项目 ID。

命令标志说明

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

vm-only

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

name
要创建的 Cloud TPU 的名称。

disk-size-gb
由 ctpu up 命令创建的虚拟机的硬盘大小（以 GB 为单位）。

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

tf-version
在虚拟机上安装的 Tensorflow ctpu 的版本。

如需详细了解 CTPU 实用程序，请参阅 CTPU 参考文档。
在出现提示时，按 y 创建 Cloud TPU 资源。
注意：第一次在项目上运行 ctpu up 时，大约需要 5 分钟来执行一些启动任务，例如 SSH 密钥传播和 API 启动。

当 ctpu up 命令执行完毕后，验证 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 数据集。

使用 ctpu 实用程序启动 Cloud TPU 资源。
```
(vm)$ ctpu up --tpu-only \
 --tpu-size=v3-8  \
 --name=resnet-tutorial \
 --zone=europe-west4-a \
 --tf-version=2.3.1
```
命令标志说明

tpu-only

创建 Cloud TPU，而不创建虚拟机。默认情况下，ctpu up 命令会同时创建虚拟机和 Cloud TPU。

tpu-size
要创建的 Cloud TPU 的类型。

name
要创建的 Cloud TPU 的名称。

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

tf-version
在虚拟机上安装的 Tensorflow ctpu 的版本。

如需详细了解 CTPU 实用程序，请参阅 CTPU 参考文档。
设置 Cloud TPU 名称变量。这可以是您在 ctpu up 中使用 --name 参数指定的名称，也可以是默认名称，即您的用户名：
```
(vm)$ export TPU_NAME=resnet-tutorial
```

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

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

导航到 ResNet-50 模型目录：
```
(vm)$ cd /usr/share/models/official/vision/image_classification/resnet/
```
注意：如果您运行的是 TensorFlow 2.1，请使用路径：/usr/share/models/official/vision/image_classification/
运行训练脚本。使用一个 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
TPU 的名称。

model_dir

用于指定在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在，此程序会自行创建相应文件夹。使用 Cloud TPU 时，model_dir 必须是 Cloud Storage 路径 (gs://...)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点，只要先前的检查点是使用相同大小的 TPU 和相同 TensorFlow 版本创建的即可。

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 上训练 ResNet 模型，请将 distribution_strategy 设置为 tpu。

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

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

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

ResNet 1 周期的训练将在 10 分钟内在 v3-8 TPU 节点上完成。训练结束时，将显示类似以下所示的输出：

I1107 20:28:57.561836 140033625347520 resnet_ctl_imagenet_main.py:222] Training 1 epochs, each epoch has 1251 steps, total steps: 1251; Eval 48 steps
I1107 20:34:09.638025 140033625347520 resnet_ctl_imagenet_main.py:358] Training loss: 0.6292637, accuracy: 0.99680257 at epoch 1
I1107 20:34:21.682796 140033625347520 resnet_ctl_imagenet_main.py:372] Test loss: 3.8977659, accuracy: 0.0% at epoch: 1
I1107 20:34:22.028973 140033625347520 resnet_ctl_imagenet_main.py:392]
Run stats:
{'train_loss': 0.6292637, 'train_acc': 0.99680257, 'eval_acc': 0.0, 'step_timestamp_log':
['BatchTimestamp <batch_index: 1, timestamp: 1573158554.11>'],
'train_finish_time': 1573158861.683073, 'eval_loss': 3.8977659>}

要训练 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: TPU 的名称。
model_dir: 用于指定在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在，此程序会自行创建。使用 Cloud TPU 时，model_dir 必须是 Cloud Storage 路径 (gs://...)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点，只要先前的检查点是使用相同大小的 TPU 和 TensorFlow 版本创建的即可。
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 上训练 ResNet 模型，请将 distribution_strategy 设置为 tpu。
log_steps: 记录时间信息（例如 examples per second）之前要执行的训练步数。
single_l2_loss_op: 计算串联权重的 L2_loss，而不是使用 Keras 每层 L2 损失。
use_tf_function: 将训练和测试步骤封装在 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 进行训练

删除为在单台设备上训练模型而创建的 Cloud TPU 资源。
```
(vm)$ ctpu delete --zone=europe-west4-a \
 --tpu-only \
 --name=resnet-tutorial
```
命令标志说明

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

tpu-only
删除 Cloud TPU。

name
要创建的 Cloud TPU 的名称。

disk-size-gb
由 ctpu up 命令创建的虚拟机的硬盘大小（以 GB 为单位）。
删除 Cloud TPU 后，创建新的 Cloud TPU Pod。运行 ctpu up 命令，并使用 tpu-size 参数指定要使用的 Pod 切片。例如，以下命令使用 v3-32 Pod 切片。
```
(vm)$ ctpu up --zone=europe-west4-a \
--tpu-only \
--name=resnet-tutorial \
--tpu-size=v3-32
```
命令标志说明

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

tpu-only

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

name
要创建的 Cloud TPU 的名称。

tpu-size
要创建的 Cloud TPU 的类型。

如需详细了解 CTPU 实用程序，请参阅 CTPU 参考文档。

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

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

导航到脚本目录：
```
(vm)$ cd /usr/share/models/official/vision/image_classification/resnet
```
注意：如果您运行的是 TensorFlow 2.1，请使用路径：/usr/share/models/official/vision/image_classification/
训练模型。
```
(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
TPU 的名称。

model_dir

用于指定在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在，此程序会自行创建。使用 Cloud TPU 时，model_dir 必须是 Cloud Storage 路径 (gs://...)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点，只要先前的检查点是使用相同大小的 Cloud TPU 和 TensorFlow 版本创建的即可。

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 上训练 ResNet 模型，请将 distribution_strategy 设置为 tpu。

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 分钟。训练和评估完成后，系统将显示类似如下内容的消息：

1107 22:45:19.821746 140317155378624 resnet_ctl_imagenet_main.py:358] Training loss: 0.22576721, accuracy: 0.838141 at epoch 1
I1107 22:45:33.892045 140317155378624 resnet_ctl_imagenet_main.py:372] Test loss: 0.26673648, accuracy: 0.0% at epoch: 1
I1107 22:45:34.851322 140317155378624 resnet_ctl_imagenet_main.py:392] Run stats:
{'train_loss': 0.22576721, 'train_acc': 0.838141, 'eval_acc': 0.0, 'step_timestamp_log': ['BatchTimestamp<batch_index: 1, timestamp: 1573166574.67>'], 'train_finish_time': 1573166733.892282, 'eval_loss': 0.26673648}

清理

为避免因本教程中使用的资源导致您的 Google Cloud Platform 帐号产生费用，请执行以下操作：

断开与 Compute Engine 实例的连接（如果您尚未这样做）：
```
(vm)$ exit
```
您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。
在 Cloud Shell 中，使用您在设置 Compute Engine 虚拟机和 Cloud TPU 时使用的 --zone 标志运行 ctpu delete。这会同时删除您的虚拟机和 Cloud TPU。
```
$ ctpu delete --zone=europe-west4-a \
  --name=resnet-tutorial
```
重要提示：如果您在运行 ctpu up 时设置了 TPU 资源名称，则必须在运行 ctpu delete 时使用 --name 标志指定该名称才能关停 TPU 资源。
运行 ctpu status 以确保未分配任何实例，避免产生不必要的 TPU 使用费。删除操作可能需要几分钟时间才能完成。如下所示的响应表明不再有已分配的实例：
```
$ ctpu status --zone=europe-west4-a
```
```
2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
No instances currently exist.
        Compute Engine VM:     --
        Cloud TPU:             --
```
如下所示运行 gsutil，将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储分区的名称：
```
$ gsutil rm -r gs://bucket-name
```

后续步骤

在本教程中，您已使用示例数据集训练 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。