使用 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 价格;并且在使用完资源之后,按照相关说明清理资源

  5. 通过在 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 下载此文件。

    # 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 下载此文件。

    # 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 价格指南,了解免费存储限制和其他价格信息。

后续步骤