使用 Cloud TPU 和 GKE 训练 Mask R-CNN

本教程介绍如何在 Cloud TPU 和 GKE 上训练 Mask RCNN 模型。

目标

创建 Cloud Storage 存储分区以保存数据集和模型输出。
创建一个 GKE 集群来管理您的 Cloud TPU 资源。
下载 Kubernetes 作业规范，其中指明了使用 TensorFlow 在 Cloud TPU 上训练 Mask RCNN 所需的资源。
在 GKE 集群中运行作业以开始训练模型。
检查日志和模型输出。

费用

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

Compute Engine
Cloud TPU
Cloud Storage

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

准备工作

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

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

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

转到项目选择器页面

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

结合使用 GKE 和 Cloud TPU 时，您的项目会使用 Google Cloud 的计费组件。如需估算您的费用，请查看 Cloud TPU 价格和 GKE 价格；并且在使用完资源之后，按照相关说明清理资源。

通过在 Cloud Console 中打开以下链接，启用本教程所需的 API：

创建 Cloud Storage 存储分区

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

转到 Cloud Console 上的 Cloud Storage 页面。

转到 Cloud Storage 页面
创建一个新的存储分区，并指定以下选项：
- 您选择的唯一名称。
- 默认存储类别：Regional
- 选择数据存储位置：region 位置：us-central1
- 选择数据的默认存储类别：标准
- 选择如何控制对象的访问权限：设置对象级层和存储分区级层权限

在 GKE 上创建集群

按照以下说明操作，使用 gcloud 命令行工具设置环境并创建一个支持 Cloud TPU 的 GKE 集群。

打开一个 Cloud Shell 窗口。

打开 Cloud Shell
配置 gcloud 命令行工具，以使用要在其中创建 Cloud TPU 的项目。
```
gcloud config set project ${PROJECT_ID}
```
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 GCP 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 TPU 访问 Cloud Storage 存储分区

您需要向 Cloud TPU 提供对 Cloud Storage 对象的读取/写入权限。为此，您必须向 Cloud TPU 使用的服务帐号授予必需的访问权限。请按照相关指南授予对存储分区的访问权限。
指定要在其中使用 Cloud TPU 资源的地区。本示例使用的是 us-central1-b 地区：
```
$ gcloud config set compute/zone us-central1-b
```
使用 gcloud container clusters 命令在 GKE 上创建一个支持 Cloud TPU 的集群。请注意，必须在提供 Cloud TPU 的地区创建 GKE 集群及其节点池，如上文中关于环境变量的部分所述。以下命令会创建一个名为 tpu-models-cluster 的集群：
```
$ gcloud container clusters create tpu-models-cluster \
--cluster-version=1.13 \
--scopes=cloud-platform \
--enable-ip-alias \
--enable-tpu \
--machine-type=n1-standard-4
```
在上述命令中：
- --cluster-version=1.13 表示该集群将使用最新的 Kubernetes 1.16 版本。您必须使用 1.13.4-gke.5 或更高版本。
- --scopes=cloud-platform 可确保集群中的所有节点都可以访问上面定义为project-name 的 Google Cloud 中的 Cloud Storage 存储分区。该集群和存储分区必须位于同一项目中。请注意，Pod 默认继承其所部署到的节点的访问权限范围。此标志向在该集群中运行的所有 Pod 授予 cloud-platform 范围。如果您想要根据每个 Pod 来限制访问权限，请参阅 GKE 指南通过服务帐号进行身份验证。
- --enable-ip-alias 表示集群使用别名 IP 范围。如需在 GKE 上使用 Cloud TPU，这是必需的设置。
- --enable-tpu 表示集群必须支持 Cloud TPU。
- 运行此作业需要 machine-type=n1-standard-4。
命令运行完毕时，将显示如下确认消息：
```
kubeconfig entry generated for tpu-models-cluster.
NAME                LOCATION       MASTER_VERSION   MASTER_IP      MACHINE_TYPE   NODE_VERSION   NUM_NODES  STATUS
tpu-models-cluster  us-central1-b  1.13.6-gke.5     35.232.204.86  n1-standard-4  1.13.6-gke.5   3          RUNNING
```

处理训练数据

您的下一个任务是创建一个作业。在 Google Kubernetes Engine 中，作业是一个表示有限任务的控制器对象。您要创建的第一个作业将下载和处理用于训练 Mask RCNN 模型的 COCO 数据集。

在 shell 环境中，创建名为 download_and_preprocess_coco_k8s.yaml 的文件，如下所示。或者，您可以从 GitHub 下载此文件。

tools/datasets/download_and_preprocess_coco_k8s.yaml

在 GitHub 上查看

# Download and preprocess the COCO dataset.
#
# Instructions:
#   1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup
#      to create a Kubernetes Engine cluster. The Job must be running at least
#      on a n1-standard-4 machine.
#   2. Change the environment variable DATA_BUCKET below to the path of the
#      Google Cloud Storage bucket where you want to store the training data.
#   3. Run `kubectl create -f download_and_preprocess_coco_k8s.yaml`.

apiVersion: batch/v1
kind: Job
metadata:
  name: download-and-preprocess-coco
spec:
  template:
    spec:
      restartPolicy: Never
      containers:
      - name: download-and-preprocess-coco
        # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile.
        image: gcr.io/tensorflow/tpu-models:r1.13
        command:
          - /bin/bash
          - -c
          - >
            DEBIAN_FRONTEND=noninteractive apt-get update &&
            cd /tensorflow_tpu_models/tools/datasets &&
            bash download_and_preprocess_coco.sh /scratch-dir &&
            gsutil -m cp /scratch-dir/*.tfrecord ${DATA_BUCKET}/coco &&
            gsutil cp /scratch-dir/raw-data/annotations/*.json ${DATA_BUCKET}/coco
        env:
          # [REQUIRED] Must specify the Google Cloud Storage location where the
          # COCO dataset will be stored.
        - name: DATA_BUCKET
          value: "gs://<my-data-bucket>/data/coco"
        volumeMounts:
        - mountPath: /scratch-dir
          name: scratch-volume
      volumes:
      - name: scratch-volume
        persistentVolumeClaim:
          claimName: scratch-disk-coco
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: scratch-disk-coco
spec:
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 100Gi

找到名为 DATA_BUCKET 的环境变量。将此变量的 value 更新为您要在 Cloud Storage 存储分区上存储 COCO 数据集的路径。例如：“gs://my-maskrcnn-bucket/coco/”

运行以下命令以在集群中创建作业：

kubectl create -f download_and_preprocess_coco_k8s.yaml

作业可能需要几分钟才能启动。您可以使用以下命令查看集群中运行的作业的状态：

kubectl get pods -w

此命令返回有关集群中运行的所有作业的信息。例如：

NAME                                 READY     STATUS      RESTARTS   AGE
download-and-preprocess-coco-xmx9s   0/1       Pending     0          10s

-w 标志指示该命令监视 pod 状态的任何更改。几分钟后，您的作业状态应更新为 Running。

如果您的作业遇到任何问题，可以将其删除，然后重试。为此，请先运行 kubectl delete -f download_and_preprocess_coco_k8s.yaml 命令，然后再次运行 kubectl create 命令。

下载和预处理完成后，您就可以运行该模型了。

运行 Mask RCNN 模型

现已全部准备就绪，您可使用 Cloud TPU 和 GKE 来运行 Mask RCNN 模型了。

在 shell 环境中，创建名为 mask_rcnn_k8s.yaml 的文件，如下所示。或者，您可以从 GitHub 下载此文件。

models/experimental/mask_rcnn/mask_rcnn_k8s.yaml

在 GitHub 上查看

# Train Mask-RCNN with Coco dataset using Cloud TPU and Google Kubernetes Engine.
#
# [Training Data]
#   Download and preprocess the COCO dataset using https://github.com/tensorflow/tpu/blob/r1.13/tools/datasets/download_and_preprocess_coco_k8s.yaml
#   if you don't already have the data.
#
# [Instructions]
#   1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup
#      to create a Kubernetes Engine cluster.
#   2. Change the environment variable MODEL_BUCKET in the Job spec to the
#      Google Cloud Storage location where you want to store the output model.
#   3. Run `kubectl create -f mask_rcnn_k8s.yaml`.

apiVersion: batch/v1
kind: Job
metadata:
  name: mask-rcnn-gke-tpu
spec:
  template:
    metadata:
      annotations:
        # The Cloud TPUs that will be created for this Job must support
        # TensorFlow 1.13. This version MUST match the TensorFlow version that
        # your model is built on.
        tf-version.cloud-tpus.google.com: "1.13"
    spec:
      restartPolicy: Never
      containers:
      - name: mask-rcnn-gke-tpu
        # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile.
        image: gcr.io/tensorflow/tpu-models:r1.13
        command:
        - /bin/sh
        - -c
        - >
            DEBIAN_FRONTEND=noninteractive apt-get update &&
            DEBIAN_FRONTEND=noninteractive apt-get install -y python-dev python-tk libsm6 libxrender1 libxrender-dev libgtk2.0-dev libxext6 libglib2.0 &&
            pip install Cython matplotlib opencv-python-headless &&
            pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' &&
            python /tensorflow_tpu_models/models/experimental/mask_rcnn/mask_rcnn_main.py
            --model_dir=${MODEL_BUCKET}
            --iterations_per_loop=500
            --config=resnet_checkpoint=${RESNET_CHECKPOINT},resnet_depth=50,use_bfloat16=true,train_batch_size=64,eval_batch_size=8,training_file_pattern=${DATA_BUCKET}/train-*,validation_file_pattern=${DATA_BUCKET}/val-*,val_json_file=${DATA_BUCKET}/instances_val2017.json,total_steps=22500
        env:
          # The Google Cloud Storage location to store dataset.
        - name: DATA_BUCKET
          value: "gs://<my-data-bucket>"
        - name: MODEL_BUCKET
          value: "gs://<my-model-bucket>/mask_rcnn"
        - name: RESNET_CHECKPOINT
          value: "gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-02-07/model.ckpt-112603"

          # Point PYTHONPATH to the top level models folder
        - name: PYTHONPATH
          value: "/tensorflow_tpu_models/models"
        resources:
          limits:
            # Request a single v3-8 Cloud TPU device to train the model.
            # A single v3-8 Cloud TPU device consists of 4 chips, each of which
            # has 2 cores, so there are 8 cores in total.
            cloud-tpus.google.com/v3: 8

找到名为 DATA_BUCKET 的环境变量。将此变量的 value 更新为您在 download_and_preprocess_coco_k8s.yaml 脚本中指定的 DATA_BUCKET 文件的路径。
找到名为 MODEL_BUCKET 的环境变量。将此变量的 value 更新为 Cloud Storage 存储分区的路径。此位置就是脚本保存训练输出的位置。如果指定的 MODEL_BUCKET 不存在，则会在您运行作业时创建。
运行以下命令以在集群中创建作业：
```
kubectl create -f mask_rcnn_k8s.yaml
```
运行此命令时，将显示以下确认消息：
```
job "mask-rcnn-gke-tpu" created
```

如处理训练数据部分所述，该作业可能需要几分钟才能启动。使用以下命令查看在您的集群中运行的作业的状态：

kubectl get pods -w

查看训练状态

您可以使用 kubectl 命令行实用程序跟踪训练状态。

运行以下命令以获取作业的状态：
```
kubectl get pods
```
在训练过程中，pod 的状态应为 RUNNING。

注意：虽然不常发生，但作业和 pod 可能会失败。如果您看到 RUNNING 以外的状态，则可能需要删除该作业并重新创建。

如需获取有关训练过程的其他信息，请运行以下命令：

$ kubectl logs job name

其中 job name 是作业的完整名称，例如 mask-rcnn-gke-tpu-abcd。

您还可以在 Cloud Console 上的 GKE 工作负载信息中心检查输出。

请注意，第一个条目需要一段时间后才会在日志中显示。您可能会看到如下内容：

I0622 18:14:31.617954 140178400511808 tf_logging.py:116] Calling model_fn.
I0622 18:14:40.449557 140178400511808 tf_logging.py:116] Create CheckpointSaverHook.
I0622 18:14:40.697138 140178400511808 tf_logging.py:116] Done calling model_fn.
I0622 18:14:44.004508 140178400511808 tf_logging.py:116] TPU job name worker
I0622 18:14:45.254548 140178400511808 tf_logging.py:116] Graph was finalized.
I0622 18:14:48.346483 140178400511808 tf_logging.py:116] Running local_init_op.
I0622 18:14:48.506665 140178400511808 tf_logging.py:116] Done running local_init_op.
I0622 18:14:49.135080 140178400511808 tf_logging.py:116] Init TPU system
I0622 18:15:00.188153 140178400511808 tf_logging.py:116] Start infeed thread controller
I0622 18:15:00.188635 140177578452736 tf_logging.py:116] Starting infeed thread controller.
I0622 18:15:00.188838 140178400511808 tf_logging.py:116] Start outfeed thread controller
I0622 18:15:00.189151 140177570060032 tf_logging.py:116] Starting outfeed thread controller.
I0622 18:15:07.316534 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:15:07.316904 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:16:13.881397 140178400511808 tf_logging.py:116] Saving checkpoints for 100 into gs://<bucket-name>/mask_rcnn/model.ckpt.
I0622 18:16:21.147114 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
I0622 18:16:21.148168 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
I0622 18:16:21.150870 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:16:21.151168 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:17:00.739207 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:17:00.739809 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:17:36.598773 140178400511808 tf_logging.py:116] global_step/sec: 2.65061
I0622 18:17:37.040504 140178400511808 tf_logging.py:116] examples/sec: 2698.56
I0622 18:17:37.041333 140178400511808 tf_logging.py:116] loss = 2.63023, step = 200 (75.893 sec)

在 gs://<bucket-name>/mask-rcnn/model.ckpt 中查看经过训练的模型。
```
gsutil ls -r gs://bucket-name/mask-rcnn/model.ckpt
```

清理

在 GKE 上使用 Cloud TPU 后，请清理资源，以避免您的 Google Cloud 帐号产生额外扣款。

如果您尚未为此会话设置项目和地区，请立即设置。请参阅本指南前面的说明。然后按照下列清理步骤进行操作：

运行以下命令以删除 GKE 集群 tpu-models-cluster，将 project-name 替换为您的 Google Cloud 项目名称：
```
$ gcloud container clusters delete tpu-models-cluster --project=project-name
```
检查完数据后，使用 gsutil 命令删除您在本教程期间创建的 Cloud Storage 存储分区。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称：
```
$ gsutil rm -r gs://bucket-name
```
请参阅 Cloud Storage 价格指南，了解免费存储限制和其他价格信息。

后续步骤

探索 TensorBoard 中的 TPU 工具。
使用以下作业规范之一运行更多模型和数据集检索作业：
- 在 GKE 上下载并预处理 COCO 数据集。
- 在 GKE 上下载并预处理 ImageNet。
- 使用 Cloud TPU 和 GKE 训练 AmoebaNet-D。
- 使用 Cloud TPU 和 GKE 训练 Inception v3。
- 使用 Cloud TPU 和 GKE 训练 RetinaNet。
使用更多 TPU 示例进行实验。