本教程介绍如何使用 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 Cloud 新手,请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源,以免产生不必要的费用。
Cloud TPU 单设备训练
本部分介绍如何为单一设备训练设置 Cloud Storage 存储桶、虚拟机和 Cloud TPU 资源。
打开一个 Cloud Shell 窗口。
为项目 ID 创建一个变量。
$ export PROJECT_ID=project-id
配置 Google Cloud CLI 以使用要在其中创建项目的项目 Cloud TPU。
$ gcloud config set project ${PROJECT_ID}
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示
Authorize Cloud Shell
页面。点击底部的Authorize
允许gcloud
使用您的凭据进行 API 调用。为 Cloud TPU 项目创建服务账号。
通过服务账号,Cloud TPU 服务可以访问其他 Google Cloud 服务。
$ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
该命令将返回以下格式的 Cloud TPU 服务账号:
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
使用以下命令创建 Cloud Storage 存储分区:
$ gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2
此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的
gcloud
命令会为您设置 TPU,还会为在上一步中设置的 Cloud TPU 服务账号设置默认权限。如果您需要更精细的权限,请查看访问级层权限。准备数据集或使用 fake_imagenet
ImageNet 是一个图片数据库。数据库中的图片被整理为 一个层次结构,层次结构中的每个节点都用数百个和 成千上万张图片。
本教程使用演示版的完整 ImageNet 数据集,该数据集又称为 fake_imagenet。此演示版本可用于测试教程,同时降低通常与使用完整 ImageNet 数据库运行模型相关的存储空间和时间要求。
此 fake_imagenet 数据集位于 Cloud Storage 上的以下位置:
gs://cloud-tpu-test-datasets/fake_imagenet
此 fake_imagenet 数据集仅用于了解如何使用 Cloud TPU 并验证端到端性能。准确率数字和保存的模型并无实际意义。
如果要使用完整的 ImageNet 数据集,请参阅 下载、预处理和上传 ImageNet 数据集。
使用
gcloud
命令启动 TPU 资源。使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需详细了解这两种虚拟机架构,请参阅系统架构。$ gcloud compute tpus tpu-vm create resnet-tutorial \ --zone=us-central2-b \ --accelerator-type=v4-8 \ --version=tpu-vm-tf-2.17.0-pjrt
命令标志说明
如需详细了解
gcloud
命令,请参阅 gcloud 参考文档。使用 SSH 连接到 TPU 虚拟机实例。连接到虚拟机后 shell 提示符从
username@projectname
更改为username@vm-name
:gcloud compute tpus tpu-vm ssh resnet-tutorial --zone=us-central2-b
设置 Cloud TPU 名称变量。
(vm)$ export TPU_NAME=local
设置 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
训练应用预期能够访问您在 Cloud Storage 中的训练数据。在训练期间,训练应用还会使用您的 Cloud Storage 存储分区来存储检查点。
创建 TPU 时,如果您将
--version
参数设置为以-pjrt
结尾的版本,请设置以下环境变量以启用 PJRT 运行时:(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
安装 TensorFlow 要求。
(vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
ResNet 训练脚本需要额外的软件包。立即安装:
(vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3
切换至存储模型的目录。
(vm)$ cd /usr/share/tpu/tensorflow/resnet50_keras
设置
PYTHONPATH
环境变量:(vm)$ /usr/share/tpu/tensorflow/resnet50_keras"
运行训练脚本。它使用 fake_imagenet 数据集并训练 ResNet 100 步。
(vm)$ resnet50.py --tpu=local --data=gs://cloud-tpu-test-datasets/fake_imagenet
命令标志说明
tpu
- TPU 的名称。
data
- 指定存储检查点和摘要的目录
模型训练期间的预测。如果指定的文件夹不存在,此程序会自行创建。使用 Cloud TPU 时,
model_dir
必须是 Cloud Storage 路径 (gs://...
)。您可以 可重复使用现有的文件夹来加载当前检查点数据并存储 其他检查点(只要先前的检查点已创建) 相同大小的 TPU 和 TensorFlow 版本。
这将训练 ResNet 100 步,并将在 大约 3 分钟。100 个步骤结束后,系统会显示类似以下所示的输出:
I0624 17:04:26.974905 140457742666816 controller.py:290] eval | step: 100 | eval time: 23.3 sec | output: {'accuracy': 0.0010141226, 'top_5_accuracy': 0.0051457332, 'validation_loss': 8.448798} eval | step: 100 | eval time: 23.3 sec | output: {'accuracy': 0.0010141226, 'top_5_accuracy': 0.0051457332, 'validation_loss': 8.448798}
您现已完成单设备训练示例。请按照以下步骤删除当前的单设备 TPU 资源。
与 Compute Engine 实例断开连接:
(vm)$ exit
您的提示符现在应为
username@projectname
,表明您位于 Cloud Shell 中。删除 TPU 资源。
$ gcloud compute tpus tpu-vm delete resnet-tutorial \ --zone=us-central2-b
命令标志说明
zone
- 您的 Cloud TPU 所在的可用区。
使用 Cloud TPU Pod 扩缩模型
在 Cloud TPU Pod 上训练模型可能需要对训练脚本进行一些更改。如需了解相关信息,请参阅在 TPU Pod 上训练。
TPU Pod 训练
本部分介绍如何为 Pod 训练设置 Cloud Storage 存储桶和 Cloud TPU 资源。
设置
打开一个 Cloud Shell 窗口。
为项目 ID 创建一个变量。
export PROJECT_ID=project-id
配置 Google Cloud CLI 以使用要在其中创建 Cloud TPU 的项目。
gcloud config set project ${PROJECT_ID}
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示
Authorize Cloud Shell
页面。点击底部的Authorize
允许gcloud
使用您的凭据进行 Google Cloud API 调用。为 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 存储分区,或使用您之前为项目创建的存储分区:
gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2
此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的
gcloud
命令会为您在上一步中设置的 Cloud TPU 服务账号设置默认权限。如果您需要更精细的权限,请查看访问级层权限。存储桶位置必须与您的 TPU 虚拟机位于同一区域。
准备数据集或使用 fake_imagenet
ImageNet 是一个图片数据库。数据库中的图像被整理为一个层次结构,该层次结构中的每个节点由成百上千个图像进行描述。
默认的 Pod 训练访问完整的 ImageNet 数据集演示版本,称为 fake_imagenet。此演示版本可让你 测试 Pod 训练,同时降低通常对存储空间和时间的要求, 与使用完整的 ImageNet 数据库训练模型相关的任务。
此 fake_imagenet 数据集仅用于了解如何使用 Cloud TPU 并验证端到端性能。准确率数字和保存的模型并无实际意义。
如果您要使用完整的 ImageNet 数据集,请参阅下载、预处理和上传 ImageNet 数据集。
创建 Cloud TPU 资源
使用
gcloud
命令启动 Cloud TPU 资源。本教程指定 v3-32 Pod。如需了解其他 Pod 选项,请参阅 TPU 版本。
$ gcloud compute tpus tpu-vm create resnet-tutorial \ --zone=us-central2-b \ --accelerator-type=4-32 \ --version=tpu-vm-tf-2.17.0-pod-pjrt
使用 SSH 连接到 TPU 虚拟机实例。连接到虚拟机后 shell 提示符从
username@projectname
更改为username@vm-name
:gcloud compute tpus tpu-vm ssh resnet-tutorial --zone=us-central2-b
导出 Cloud TPU 设置变量:
(vm)$ export ZONE=us-central2-b (vm)$ export STORAGE_BUCKET=gs://bucket-name (vm)$ export TPU_NAME=resnet-tutorial (vm)$ export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x-pod
训练应用预期能够访问您在 Cloud Storage 中的训练数据。在训练期间,训练应用还会使用您的 Cloud Storage 存储分区来存储检查点。
ResNet 训练脚本需要额外的软件包。立即安装:
(vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3
安装 TensorFlow 要求。{: id='setup-env'}:
(vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
设置
PYTHONPATH
环境变量:(vm)$ export PYTHONPATH="PYTHONPATH=/usr/share/tpu/tensorflow/resnet50_keras" (vm)$ export TPU_LOAD_LIBRARY=0
切换至存储模型的目录。
(vm)$ cd /usr/share/tpu/tensorflow/resnet50_keras
训练模型。
(vm)$ resnet50.py --tpu=${TPU_NAME} --data=gs://cloud-tpu-test-datasets/fake_imagenet
命令标志说明
tpu
- TPU 的名称。
data
- 用于指定在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在,此程序会自行创建。
使用 Cloud TPU 时,
model_dir
必须是 Cloud Storage 路径 (gs://...
)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点,只要先前的检查点是使用相同大小的 Cloud TPU 和相同 TensorFlow 版本创建的即可。
此过程使用 fake_imagenet 数据集对该模型训练 100 次训练 和 13 个评估步骤。本次培训需要大约 2 分钟的时间, v3-32 Cloud TPU。训练和评估完成后, 会显示类似如下内容的消息:
{'accuracy': 0.0009716797, 'learning_rate': 0.10256411, 'top_5_accuracy': 0.0049560545, 'training_loss': 8.5587225} train | step: 100 | steps/sec: 1.2 | output: {'accuracy': 0.0009716797, 'learning_rate': 0.10256411, 'top_5_accuracy': 0.0049560545, 'training_loss': 8.5587225} eval | step: 100 | eval time: 24.8 sec | output: {'accuracy': 0.0010141226, 'top_5_accuracy': 0.004356971, 'validation_loss': 8.50038} eval | step: 100 | eval time: 24.8 sec | output: {'accuracy': 0.0010141226, 'top_5_accuracy': 0.004356971, 'validation_loss': 8.50038}
清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。
断开与 Compute Engine 实例的连接(如果您尚未这样做):
(vm)$ exit
您的提示符现在应为
username@projectname
,表明您位于 Cloud Shell 中。删除您的 Cloud TPU 和 Compute Engine 资源。
$ gcloud compute tpus tpu-vm delete resnet-tutorial \ --zone=us-central2-b
通过运行
gcloud compute tpus tpu-vm list
验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出不应包含本教程中创建的任何 TPU 资源:$ gcloud compute tpus tpu-vm list --zone=us-central2-b
删除您为本教程创建的 Cloud Storage 存储桶:
$ gcloud storage rm gs://bucket-name --recursive
后续步骤
TensorFlow Cloud TPU 教程通常使用示例数据集来训练模型。此训练的结果不能用于推理。接收者 使用模型进行推理,就可以使用公开可用的 或您自己的数据集。在 Cloud TPU 上训练的 TensorFlow 模型通常要求数据集采用 TFRecord 格式。
您可以使用数据集转换工具示例将图片分类数据集转换为 TFRecord 格式。如果您未使用图片分类模型,则必须自行将数据集转换为 TFRecord 格式。如需了解详情,请参阅 TFRecord 和 tf.Example。
超参数调节
要使用数据集提升模型的性能,您可以调整模型的 超参数。您可以找到所有用户共有的超参数的相关信息 支持 TPU 的 GitHub。 您可以在每个模型的源代码中寻找模型专用超参数的相关信息。如需详细了解超参数调优,请参阅概览 超参数调优 和 Tune 超参数。
推断
训练模型后,即可将其用于推理(也称为 预测)。您可以使用 Cloud TPU 推理转换器工具准备和优化 TensorFlow 模型,以便在 Cloud TPU v5e 上进行推理。如需详细了解 Cloud TPU v5e 上的推理,请参阅 Cloud TPU v5e 推理简介。
- 遵循数据集转换教程,了解如何使用自己的数据替代 fake_imagenet 或 ImageNet 数据集来进行训练和评估。本教程介绍了如何使用图片分类数据转换器 用于将用于图片分类的原始数据集转换为 可供 Cloud TPU TensorFlow 模型使用的 TFRecord。
- 运行 Cloud TPU Colab 展示了如何使用自己的资源运行图片分类模型 图片数据。
- 浏览其他 Cloud TPU 教程。
- 了解如何使用 TensorBoard 中的 TPU 监控工具。
- 了解如何使用 Cloud TPU 和 GKE 训练 ResNet。