本文档演示如何使用 Cloud TPU 和 COCO 数据集运行 ShapeMask 模型。
以下说明假设您已熟悉如何在 Cloud TPU 上运行模型。如果您刚接触 Cloud TPU，则可以参阅快速入门获取基本介绍。
如果您计划在 TPU Pod 切片上训练，请参阅在 TPU Pod 上训练 了解 Pod 切片所需的参数更改。
目标
- 准备 COCO 数据集
- 创建 Cloud Storage 存储分区以保存数据集和模型输出
- 设置 TPU 资源以进行训练和评估
- 在单个 Cloud TPU 或 Cloud TPU Pod 上运行训练和评估
费用
在本文档中，您将使用 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 资源。
如果您计划在 TPU Pod 切片上训练，请参阅在 TPU Pod 上训练，以了解在 Pod 切片上训练所需的更改。
在 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 凭据进行 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
准备 COCO 数据集
本教程使用 COCO 数据集。Cloud Storage 存储桶上的数据集需要采用 TFRecord 格式以用于训练。
存储桶位置必须与虚拟机 (VM) 和 TPU 节点位于同一区域。虚拟机和 TPU 节点位于特定可用区中， 它们是某个区域内的细分
Cloud Storage 存储桶存储您用于训练模型的数据
训练结果。本教程中使用的 gcloud compute tpus tpu-vm 工具可设置
您在 Cloud TPU 和 Cloud TPU 中
上一步。如果您需要更精细的权限，请查看访问级层权限。
如果您已在 Cloud Storage 存储桶中 位于您将用来执行以下操作的可用区中 训练模型时，您可以启动 TPU 资源并 准备 Cloud TPU 进行训练。否则，请按照以下步骤准备数据集。
在 Cloud Shell 中，使用您的项目 ID 配置
gcloud。export PROJECT_ID=project-id gcloud config set project ${PROJECT_ID}
在 Cloud Shell 中，使用以下命令创建 Cloud Storage 存储桶：
gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2
创建 Compute Engine 虚拟机以下载和预处理数据集。有关 请参阅 创建并启动 Compute Engine 实例。
$ gcloud compute instances create vm-name \ --zone=us-central2-b \ --image-family=ubuntu-2204-lts \ --image-project=ubuntu-os-cloud \ --machine-type=n1-standard-16 \ --boot-disk-size=300GB
使用 SSH 连接到 Compute Engine 虚拟机：
$ gcloud compute ssh vm-name --zone=us-central2-b
连接到虚拟机时，Shell 提示符会从
username@projectname至username@vm-name。设置两个变量，一个用于先前创建的存储桶，另一个用于保存存储桶中的训练数据 (
DATA_DIR) 的目录。(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco
安装预处理数据所需的软件包。
(vm)$ sudo apt-get update && \ sudo apt-get install python3-pip && \ sudo apt-get install -y python3-tk && \ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow numpy absl-py tensorflow && \ pip3 install --user "git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI"
运行
download_and_preprocess_coco.sh脚本以进行转换 将 COCO 数据集转换为一组 TFRecord 文件 (*.tfrecord)， 符合预期。(vm)$ git clone https://github.com/tensorflow/tpu.git (vm)$ sudo bash tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco
这会安装所需的库，然后运行预处理脚本。它会在您的本地数据目录中输出
*.tfrecord文件。COCO 下载和转换脚本大约需要 1 小时才能完成。将数据复制到您的 Cloud Storage 存储桶。
将数据转换为 TFRecord 格式后，使用 gcloud CLI 将数据从本地存储空间复制到 Cloud Storage 存储桶。您必须 还可以复制注解文件这些文件有助于验证模型的 性能
(vm)$ gcloud storage cp ./data/dir/coco/*.tfrecord ${DATA_DIR} (vm)$ gcloud storage cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}
与 Compute Engine 虚拟机断开连接：
(vm)$ exit
您的提示符现在应为
username@projectname，表明您位于 Cloud Shell 中。删除您的 Compute Engine 虚拟机：
$ gcloud compute instances delete vm-name \ --zone=us-central2-b
启动 TPU 资源并训练模型
使用
gcloud命令启动 TPU 资源。$ gcloud compute tpus tpu-vm create shapemask-tutorial \ --zone=europe-west4-a \ --accelerator-type=v3-8 \ --version=tpu-vm-tf-2.17.0-pjrt
命令标志说明
如需详细了解
gcloud命令，请参阅 gcloud 参考文档。使用 SSH 连接到 Compute Engine 实例。连接到网络后 您的 Shell 提示符会从
username@projectname更改为username@vm-name:gcloud compute tpus tpu-vm ssh shapemask-tutorial --zone=europe-west4-a
安装 TensorFlow 要求。
(vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
训练脚本需要额外的软件包。立即安装：
(vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3
设置存储桶名称变量。将 bucket-name 替换为存储桶的名称：
(vm)$ export STORAGE_BUCKET=gs://bucket-name
设置 Cloud TPU 名称变量。
(vm)$ export TPU_NAME=local
设置
PYTHONPATH环境变量：(vm)$ export PYTHONPATH="/usr/share/tpu/models:${PYTHONPATH}"
切换至存储模型的目录。
(vm)$ cd /usr/share/tpu/models/official/legacy/detection
添加一些必需的环境变量：
(vm)$ export RESNET_CHECKPOINT=gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07 (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco (vm)$ export TRAIN_FILE_PATTERN=${DATA_DIR}/train-* (vm)$ export EVAL_FILE_PATTERN=${DATA_DIR}/val-* (vm)$ export VAL_JSON_FILE=${DATA_DIR}/instances_val2017.json (vm)$ export SHAPE_PRIOR_PATH=gs://cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/shapemask
创建 TPU 时，如果您将
--version参数设置为以-pjrt时，请设置以下环境变量以启用 PJRT 运行时：(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
训练 ShapeMask 模型：
以下脚本在 v3-8 TPU 上运行一个示例训练，其中将仅训练 100 个步骤，完成大约需要 10 分钟。在 v3-8 TPU 上，训练到收敛需要大约 22500 个步骤和大约 6 小时。
(vm)$ python3 main.py \ --strategy_type=tpu \ --tpu=${TPU_NAME} \ --model_dir=${MODEL_DIR} \ --mode=train \ --model=shapemask \ --params_override="{train: {total_steps: 100, learning_rate: {init_learning_rate: 0.08, learning_rate_levels: [0.008, 0.0008], learning_rate_steps: [15000, 20000], }, checkpoint: { path: ${RESNET_CHECKPOINT},prefix: resnet50}, train_file_pattern: ${TRAIN_FILE_PATTERN}}, shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}, shapemask_parser: {output_size: [640, 640]}}"
命令标志说明
strategy_type- 如需在 TPU 上训练 Shapemask 模型，您必须设置
将
distribution_strategy更改为tpu。 tpu- Cloud TPU 的名称。您可以使用
TPU_NAME环境变量。 model_dir- 在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在，此程序会自行创建。
使用 Cloud TPU 时，
model_dir必须 为 Cloud Storage 路径 (gs://...)。您可以重复使用 一个现有的文件夹，用于加载当前检查点数据并存储 其他检查点，只要先前的检查点 使用相同大小的 Cloud TPU 和 TensorFlow 版本。 mode- 将其设置为
train以训练模型，或设置为eval来评估模型。 params_override- 一个 JSON 字符串，会替换默认脚本参数。如需详细了解脚本参数，请参阅
/usr/share/models/official/legacy/detection/main.py。
训练完成后，系统将显示如下所示的消息：
Train Step: 100/100 / loss = {'total_loss': 10.815635681152344, 'loss': 10.815635681152344, 'retinanet_cls_loss': 1.4915691614151, 'l2_regularization_loss': 4.483549118041992, 'retinanet_box_loss': 0.013074751943349838, 'shapemask_prior_loss': 0.17314358055591583, 'shapemask_coarse_mask_loss': 1.953366756439209, 'shapemask_fine_mask_loss': 2.216097831726074, 'model_loss': 6.332086086273193, 'learning_rate': 0.021359999} / training metric = {'total_loss': 10.815635681152344, 'loss': 10.815635681152344, 'retinanet_cls_loss': 1.4915691614151, 'l2_regularization_loss': 4.483549118041992, 'retinanet_box_loss': 0.013074751943349838, 'shapemask_prior_loss': 0.17314358055591583, 'shapemask_coarse_mask_loss': 1.953366756439209, 'shapemask_fine_mask_loss': 2.216097831726074, 'model_loss': 6.332086086273193, 'learning_rate': 0.021359999}运行脚本以评估 ShapeMask 模型。此过程在 v3-8 TPU 上大约需要 10 分钟：
(vm)$ python3 main.py \ --strategy_type=tpu \ --tpu=${TPU_NAME} \ --model_dir=${MODEL_DIR} \ --checkpoint_path=${MODEL_DIR} \ --mode=eval_once \ --model=shapemask \ --params_override="{eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 }, shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}, shapemask_parser: {output_size: [640, 640]}}"
命令标志说明
strategy_type- 如需在 TPU 上训练 Shapemask 模型，您必须设置
将
distribution_strategy更改为tpu。 tpu- Cloud TPU 的名称。您可以使用
TPU_NAME环境变量。 model_dir- 在模型运行期间存储检查点和摘要的目录
训练。如果指定的文件夹不存在，此程序会自行创建。使用
Cloud TPU，
model_dir必须为 Cloud Storage 路径 (gs://...)。您可以重复使用 以加载当前检查点数据和存储其他 检查点，前提是先前的检查点是使用 大小和 TensorFlow 版本相同的 Cloud TPU。 mode- 将此值设置为
train以训练模型，或设置为eval以评估模型。 params_override- 一个 JSON 字符串，会替换默认脚本参数。如需详细了解脚本参数，请参阅
/usr/share/models/official/legacy/detection/main.py。
评估完成后，系统将显示类似如下内容的消息：
DONE (t=5.47s). Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.000 Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.000 Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.000 Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000 Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000 Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000 Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000
您现已完成单设备训练和评估。使用 删除当前的单设备 TPU 资源的步骤。
与 Compute Engine 实例断开连接：
(vm)$ exit
您的提示符现在应为
username@projectname，表明您位于 Cloud Shell 中。删除 TPU 资源。
$ gcloud compute tpus tpu-vm delete shapemask-tutorial \ --zone=europe-west4-a
命令标志说明
zone- 您的 Cloud TPU 所在的地区 实际居住地
此时，您可以结束本教程并清理，也可以继续并探索在 Cloud TPU Pod 上运行模型。
使用 Cloud TPU Pod 扩缩模型
在 Cloud TPU Pod 上训练模型可能需要对训练脚本进行一些更改。如需了解相关信息，请参阅在 TPU Pod 上训练。
TPU Pod 训练
打开一个 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使用您的 凭据。为 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=europe-west4
如果您之前已准备的 COCO 数据集并将其移至存储桶，则可以再次使用它来进行 Pod 训练。如果您尚未准备好 COCO 数据集，请立即准备并返回此处设置 Pod 训练。
启动 Cloud TPU Pod
本教程指定 v3-32 Pod。如需了解其他 Pod 选项，请参阅 TPU 版本。
$ gcloud compute tpus tpu-vm create shapemask-tutorial \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.17.0-pod-pjrt
命令标志说明
--project使用 SSH 连接到 Compute Engine 实例。连接到网络后 您的 Shell 提示符会从
username@projectname更改为username@vm-name:gcloud compute tpus tpu-vm ssh shapemask-tutorial --zone=europe-west4-a
安装 TensorFlow 要求。
(vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
训练脚本需要额外的软件包。立即安装：
(vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3
设置以下环境变量，将 bucket-name 替换为 Cloud Storage 存储分区的名称：
(vm)$ export STORAGE_BUCKET=gs://bucket-name
训练应用应该能够访问 Cloud Storage 中的训练数据。在训练期间，训练应用还会使用您的 Cloud Storage 存储分区来存储检查点。
更新所需的训练变量。
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/shapemask-pods (vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco (vm)$ export RESNET_CHECKPOINT=gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07 (vm)$ export TRAIN_FILE_PATTERN=${DATA_DIR}/train-* (vm)$ export EVAL_FILE_PATTERN=${DATA_DIR}/val-* (vm)$ export VAL_JSON_FILE=${DATA_DIR}/instances_val2017.json (vm)$ export SHAPE_PRIOR_PATH=gs://cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy
设置一些所需的环境变量：
(vm)$ export PYTHONPATH="/usr/share/tpu/models:${PYTHONPATH}" (vm)$ export TPU_LOAD_LIBRARY=0
切换至存储模型的目录。
(vm)$ cd /usr/share/tpu/models/official/legacy/detection
启动 Pod 训练。
示例训练会运行 20 个步骤，大约需要 10 分钟 在 v3-32 TPU 节点上完成。要训练收敛，大约需要 在 v3-32 TPU Pod 上完成训练大约需要 2 小时。
(vm)$ python3 main.py \ --strategy_type=tpu \ --tpu=${TPU_NAME} \ --model_dir=${MODEL_DIR} \ --mode=train \ --model=shapemask \ --params_override="{train: { batch_size: 128, iterations_per_loop: 500, total_steps: 20, learning_rate: {'learning_rate_levels': [0.008, 0.0008], 'learning_rate_steps': [10000, 13000] }, checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}}, shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}} }"
命令标志说明
strategy_type- 如需在 TPU 上训练 Shapemask 模型，您必须将
distribution_strategy设置为tpu。 tpu- Cloud TPU 的名称。这是使用
TPU_NAME环境变量设置的。 model_dir- 在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在，此程序会自行创建。时间
使用 Cloud TPU 时，
model_dir必须是 Cloud Storage 路径 (gs://...)。您可以重复使用 以加载当前检查点数据和存储其他 检查点，前提是先前的检查点是使用 大小和 TensorFlow 版本相同的 Cloud TPU。 mode- 将此值设置为
train以训练模型，或设置为eval以评估模型。 params_override- 一个 JSON 字符串，会替换默认脚本参数。如需详细了解脚本参数，请参阅
/usr/share/models/official/legacy/detection/main.py。
清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。
断开与 Compute Engine 实例的连接（如果您尚未这样做）：
(vm)$ exit
您的提示符现在应为
username@projectname，表明您位于 Cloud Shell 中。删除您的 Cloud TPU 和 Compute Engine 资源。
$ gcloud compute tpus tpu-vm delete shapemask-tutorial \ --zone=europe-west4-a
通过运行
gcloud compute tpus tpu-vm list验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出 不应包含本教程中创建的任何 TPU 资源：$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
如下所示运行 gcloud CLI，将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储桶的名称：
$ gcloud storage rm gs://bucket-name --recursive
后续步骤
使用其他图片大小进行训练
您可以使用更大的神经网络（例如 ResNet-101， ）。较大的输入图片和更强大的神经网络将生成耗时更长但更精确的模型。
使用其他基础
或者，您也可以探索如何使用自己的数据集预训练 ResNet 模型 并将其用作 ShapeMask 模型的基础。再进行一些工作，您就可以用新的神经网络来替换 ResNet。最后，如果您有兴趣实现自己的对象检测模型，则可以基于该骨干网进行进一步的实验。