使用 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
  • 云端存储

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

准备工作

  1. 登录您的 Google 帐号。

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

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

    转到项目选择器页面

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

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

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

创建 Cloud Storage 存储分区

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

  1. 转到 Cloud Console 上的 Cloud Storage 页面。

    转到 Cloud Storage 页面

  2. 创建一个新的存储分区，并指定以下选项：

    • 您选择的唯一名称。
    • 默认存储类别：Regional
    • 选择数据存储位置：region 位置：us-central1
    • 选择数据的默认存储类别：标准
    • 选择如何控制对象的访问权限：设置对象级层和存储分区级层权限

授权 Cloud TPU 访问 Cloud Storage 存储分区

您需要向 Cloud TPU 提供对 Cloud Storage 对象的读取/写入权限。为此，您必须向 Cloud TPU 使用的服务帐号授予必需的访问权限。请按照相关指南授予对存储分区的访问权限。

在 GKE 上创建集群

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

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

  2. 指定您的 Google Cloud 项目：

    $ gcloud config set project project-name

    其中 project-name 是您的 Google Cloud 项目的名称。

  3. 指定要在其中使用 Cloud TPU 资源的地区。 本示例使用的是 us-central1-b 地区：

    $ gcloud config set compute/zone us-central1-b

  4. 使用 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 数据集。

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

    tools/datasets/download_and_preprocess_coco_k8s.yaml 
    # 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

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

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

    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 模型了。

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

    models/experimental/mask_rcnn/mask_rcnn_k8s.yaml 
    # 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

  2. 找到名为 DATA_BUCKET 的环境变量。将此变量的 value 更新为您在 download_and_preprocess_coco_k8s.yaml 脚本中指定的 DATA_BUCKET 文件的路径。

  3. 找到名为 MODEL_BUCKET 的环境变量。将此变量的 value 更新为 Cloud Storage 存储分区的路径。此位置就是脚本保存训练输出的位置。如果指定的 MODEL_BUCKET 不存在，则会在您运行作业时创建。

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

    kubectl create -f mask_rcnn_k8s.yaml

    运行此命令时，将显示以下确认消息：

    job "mask-rcnn-gke-tpu" created

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

kubectl get pods -w

查看训练状态

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

  1. 运行以下命令以获取作业的状态：

    kubectl get pods

    在训练过程中，pod 的状态应为 RUNNING。

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

    $ 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)

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

    gsutil ls -r gs://bucket-name/mask-rcnn/model.ckpt

清理

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

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

  1. 运行以下命令以删除 GKE 集群 tpu-models-cluster，将 project-name 替换为您的 Google Cloud 项目名称：

    $ gcloud container clusters delete tpu-models-cluster --project=project-name

  2. 检查完数据后，使用 gsutil 命令删除您在本教程期间创建的 Cloud Storage 存储分区。将 bucket-name 替换为您的 Cloud Storage 存储分区的名称：

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

    请参阅 Cloud Storage 价格指南，了解免费存储限制和其他价格信息。

后续步骤