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

本文档介绍了 RetinaNet 对象检测模型的实现。您可以在 GitHub 上找到相关代码。

以下说明假设您已熟悉如何在 Cloud TPU 上运行模型。如果您刚接触 Cloud TPU,则可以参阅快速入门获取基本介绍。

如果您计划在 TPU Pod 切片上训练,请参阅在 TPU Pod 上训练,以了解 Pod 切片所需的参数更改。

目标

  • 创建 Cloud Storage 存储分区以保存数据集和模型输出
  • 准备 COCO 数据集
  • 设置 Compute Engine 虚拟机和 Cloud TPU 节点以进行训练和评估
  • 在单个 Cloud TPU 或 Cloud TPU Pod 上运行训练和评估

费用

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

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

准备工作

本部分介绍如何设置 Cloud Storage 存储分区和 Compute Engine 虚拟机。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. 配置 gcloud 命令行工具,以使用要在其中创建 Cloud TPU 的项目。

    gcloud config set project ${PROJECT_ID}
    
  4. 为 Cloud TPU 项目创建服务帐号。

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

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

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

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

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

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

  6. 使用 ctpu up 命令启动 Compute Engine 虚拟机。

    $ ctpu up --vm-only \
    --name=retinanet-tutorial \
    --disk-size-gb=300 \
    --machine-type=n1-standard-8 \
    --zone=europe-west4-a \
    --tf-version=2.3.1
    

    命令标志说明

    vm-only
    创建虚拟机,而不创建 Cloud TPU。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    name
    要创建的 Cloud TPU 的名称。
    disk-size-gb
    ctpu up 命令创建的虚拟机的硬盘大小(以 GB 为单位)。
    machine-type
    要创建的 Compute Engine 虚拟机的机器类型
    zone
    拟在其中创建 Cloud TPU 的区域
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。

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

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

    ctpu up 命令执行完毕后,验证 shell 提示符已从 username@projectname 更改为 username@vm-name。此变化表明您现已登录 Compute Engine 虚拟机。

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

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

  8. 安装另外的软件包

    RetinaNet 训练应用需要几个额外的软件包。 请立即安装:

    (vm)$ sudo apt-get install -y python3-tk
    (vm)$ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow
    
    (vm)$ pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI'
    
    (vm)$ sudo pip3 install --user -r /usr/share/models/official/requirements.txt
    

准备 COCO 数据集

COCO 数据集将存储在您的 Cloud Storage 中,因此请设置一个存储分区变量来指定您创建的存储分区名称:

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/coco

运行 download_and_preprocess_coco.sh 脚本,将 COCO 数据集转换为训练应用所需的一组 TFRecord (*.tfrecord)。

(vm)$ sudo bash /usr/share/tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

这会安装所需的库,然后运行预处理脚本。它会在您的本地数据目录中输出很多 *.tfrecord 文件。COCO 下载和转换脚本大约需要 1 小时才能完成。

将数据复制到 Cloud Storage 存储分区

将数据转换为 TFRecord 后,使用 gsutil 命令将其从本地存储空间复制到 Cloud Storage 存储分区。您还必须复制注释文件。这些文件有助于验证模型的性能。

(vm)$ gsutil -m cp ./data/dir/coco/*.tfrecord ${DATA_DIR}
(vm)$ gsutil cp ./data/dir/coco/raw-data/annotations/*.json ${DATA_DIR}

设置训练环境

  1. 使用 ctpu up 命令创建 Cloud TPU。

    (vm)$ ctpu up --tpu-only \
      --tpu-size=v3-8 \
      --zone=europe-west4-a \
      --name=retinanet-tutorial \
      --tf-version=2.3.1
    

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

    命令标志说明

    tpu-only
    仅创建 Cloud TPU。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    tpu-size
    要创建的 Cloud TPU 的类型
    zone
    拟在其中创建 Cloud TPU 的区域
    name
    要创建的 Cloud TPU 的名称。
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。
    参数 说明

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

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

    ctpu up 命令执行完毕后,验证 shell 提示符已从 username@projectname 更改为 username@vm-name。 此变化表明您现已登录 Compute Engine 虚拟机。

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

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

  3. 为 TPU 名称创建环境变量。

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. 使用以下命令将顶级 /models 文件夹添加到 Python 路径:

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    

单个 Cloud TPU 设备训练

以下训练脚本在 Cloud TPU v3-8 上运行。 您也可以在 Cloud TPU v2-8 上运行,但这将需要更长的时间。

以下示例脚本仅有 10 个训练步骤,并且在 v3-8 TPU 节点上的运行时间不超过 5 分钟。在 Cloud TPU v3-8 TPU 上,训练到收敛大约需要 22,500 个步骤和大约 1 1/2 个小时。

  1. 设置以下环境变量:

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-train
    (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
    
  2. 运行训练脚本:

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
         --strategy_type=tpu \
         --tpu=${TPU_NAME} \
         --model_dir=${MODEL_DIR} \
         --mode="train" \
         --params_override="{ type: retinanet, train: { total_steps: 10, 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}, eval_samples: 5000 } }"
    

    命令标志说明

    strategy_type
    要在 TPU 上训练 RetinaNet 模型,您必须将 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/vision/detection/main.py

在训练模型时,您可以通过日志输出来查看进度。 与以下内容类似的输出表明训练正常进行:

31517803669, 'learning_rate': 0.08, 'box_loss': 0.0006472870009019971,
'l2_regularization_loss': 0.09328179806470871}
I0210 21:59:19.888985 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_2500.ckpt-5
I0210 22:01:07.714749 139927795508992 distributed_executor.py:446]
Train Step: 3000/22500  / loss = {'model_loss': 0.08362223953008652,
'total_loss': 0.17120523750782013, 'cls_loss': 0.057121846824884415,
'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361} / training metric =
{'model_loss': 0.08362223953008652, 'total_loss': 0.17120523750782013, 'cls_loss': 0.0
57121846824884415, 'learning_rate': 0.08, 'box_loss': 0.0005300078773871064,
'l2_regularization_loss': 0.08758299797773361}
I0210 22:01:15.813422 139927795508992 distributed_executor.py:49]
Saving model as TF checkpoint: gs://bucket-eu/retinanet-model/ctl_step_3000.ckpt-6

单个 Cloud TPU 设备评估

以下过程使用 COCO 评估数据。运行评估步骤大约需要 10 分钟。

  1. 设置以下环境变量:

    (vm)$ export EVAL_SAMPLES=5000
    
  2. 运行评估脚本:

    (vm)$ python3 /usr/share/models/official/vision/detection/main.py \
          --strategy_type=tpu \
          --tpu=${TPU_NAME} \
          --model_dir=${MODEL_DIR} \
          --mode="eval" \
          --params_override="{ type: retinanet, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: ${EVAL_SAMPLES} } }"
    

    命令标志说明

    strategy_type
    要在 TPU 上训练 RetinaNet 模型,您必须将 distribution_strategy 设置为 tpu
    tpu
    Cloud TPU 的名称。这是使用 TPU_NAME 环境变量设置的。
    model_dir
    在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在,此程序会自行创建。 使用 Cloud TPU 时,model_dir 必须是 Cloud Storage 路径 (gs://...)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点,只要先前的检查点是使用相同大小的 Cloud TPU 和相同 TensorFlow 版本创建的即可。
    mode
    model 设置为 eval 以评估模型。
    params_override
    一个 JSON 字符串,会替换默认脚本参数。如需详细了解脚本参数,请参阅 /usr/share/models/official/vision/detection/main.py

    在评估结束时,您会在控制台中看到类似如下内容的信息:

    Accumulating evaluation results...
    DONE (t=7.66s).
     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
    

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

使用 Cloud TPU Pod 扩缩模型

您可以使用 Cloud TPU Pod 扩缩模型,以便更快获得结果。完全受支持的 RetinaNet 模型可以与 v2-32 Pod 切片搭配使用:

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

    (vm)$ ctpu delete --tpu-only --zone=europe-west4-a --name=retinanet-tutorial
  2. 运行 ctpu up 命令,并使用 tpu-size 参数指定要使用的 Pod 切片。例如,以下命令使用 v2-32 Pod 切片。

    (vm)$ ctpu up --tpu-only \
    --tpu-size=v2-32  \
    --zone=europe-west4-a \
    --name=retinanet-tutorial \
    --tf-version=2.3.1

    命令标志说明

    tpu-only
    仅创建 Cloud TPU。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    tpu-size
    要创建的 Cloud TPU 的类型
    zone
    拟在其中创建 Cloud TPU 的区域
    name
    要创建的 Cloud TPU 的名称。
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。

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

  3. 设置 Cloud TPU 名称变量。这可以是您使用 --name 参数设置的名称,也可以是默认名称,或者您的用户名:

    (vm)$ export TPU_NAME=retinanet-tutorial
    
  4. 设置以下环境变量:

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-pod
    
  5. 在 v2-32 TPU 节点上运行 Pod 训练脚本

    以下示例训练脚本是在 Cloud TPU v2-32 Pod 上运行的。它只有 10 个训练步骤,运行时间不超过 5 分钟。在 v2-32 TPU Pod 上,训练到收敛大约需要 2109 个步骤和大约 50 分钟。

    (vm)$  python3 /usr/share/models/official/vision/detection/main.py \
        --strategy_type=tpu \
        --tpu=${TPU_NAME} \
        --model_dir=${MODEL_DIR} \
        --mode="train" \
        --params_override="{ type: retinanet, train: { total_steps: 10, batch_size: 256, 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}, eval_samples: 5000 } }"
    

    命令标志说明

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

清理

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

  1. 与 Compute Engine 虚拟机断开连接:

    (vm)$ exit
    

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

  2. 在虚拟机或 Cloud Shell 中,使用以下命令删除虚拟机和 Cloud TPU:

    $ ctpu delete --name=retinanet-tutorial \
      --zone=europe-west4-a
    
  3. 运行 ctpu status 指定您的区域,以确保您未分配任何实例,从而避免产生不必要的 TPU 使用费。删除操作可能需要几分钟时间才能完成。如下所示的响应表明不再有已分配的实例:

    $ ctpu status --name=retinanet-tutorial \
      --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:             --
    
  4. 如下所示运行 gsutil,将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储分区的名称:

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

后续步骤

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

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

超参数调节

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

推理

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

使用其他图像大小进行训练

您可以使用更大的骨干网(例如 ResNet-101,而不是 ResNet-50)进行探索。更大的输入图像和更强大的骨干网将生成耗时更长但更精确的模型。

使用其他基础

或者,您也可以探索如何利用自己的数据集预训练 ResNet 模型,并以此为基础训练 RetinaNet 模型。再多进行一些工作,您还可以使用替代骨干网来替换掉 ResNet。最后,如果您有兴趣实现自己的对象检测模型,则可以基于该骨干网进行进一步的实验。