使用 GKE 在 Cloud TPU 上训练 ResNet

目标

本教程介绍如何在 Cloud TPU 和 GKE 上训练 TensorFlow ResNet-50 模型。

本教程将使用专为测试目的而提供的虚构数据集,引导您完成下列步骤来训练模型:

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

要求和限制

定义配置时,请注意以下事项:

  • 您必须使用 GKE 1.13.4-gke.5 或更高版本。您可以按照如下所述在 gcloud container clusters create 命令中添加 --cluster-version 参数以指定版本。要详细了解版本信息,请参阅 SDK 文档
  • 您必须使用 TensorFlow 1.15.3 或更高版本。您应该按如下所述在 Kubernetes Pod 规范中指定 Cloud TPU 所用的 TensorFlow 版本。
  • 您必须在提供 Cloud TPU 的地区创建 GKE 集群和节点池,还必须创建 Cloud Storage 存储分区,将您的训练数据和模型保存在 GKE 集群所属的区域中。 本教程使用 us-central1-b 地区。
  • 每个容器最多可以申请一个 Cloud TPU,但 Kubernetes Pod 中的多个容器可以各自申请一个 Cloud TPU。
  • 集群自动扩缩程序在 GKE 1.11.4-gke.12 及更高版本上支持 Cloud TPU。
  • 您的 GKE 集群必须是 VPC 原生集群

准备工作

  1. 登录您的 Google Cloud 帐号。如果您是 Google Cloud 新手,请创建一个帐号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  2. 在 Google Cloud Console 的项目选择器页面上,选择或创建一个 Google Cloud 项目。

    转到“项目选择器”

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

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

  5. 在 Cloud Console中启用以下API:

创建服务帐号和 Cloud Storage 存储分区

您需要一个 Cloud Storage 存储分区来存储机器学习模型的训练结果。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. 为 Cloud TPU 项目创建服务帐号。

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

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

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    
  4. 转到 Cloud Console 上的 Cloud Storage 页面。

    转到 Cloud Storage 页面

  5. 创建一个新的存储分区,并指定以下选项:

    • 您选择的唯一名称。
    • 位置类型:region
    • 位置:us-central1
    • 默认存储类别:Standard
    • 访问权限控制机制:fine-grained

    在使用存储分区之前,您需要向 Cloud TPU 服务帐号授予存储分区的访问权限。使用上面的访问权限控制机制链接为您的 Cloud TPU 服务帐号设置精细 ACL。

创建支持 Cloud TPU 的 GKE 集群

您可以在新的 GKE 集群上启用 Cloud TPU 支持。

创建支持 Cloud TPU 的新集群

您可以使用 Cloud Console 或 gcloud 工具创建支持 Cloud TPU 的集群。

要查看相关说明,请选择下述某一选项:

控制台

按照下述说明创建一个支持 Cloud TPU 的 GKE 集群:

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

    转到 GKE 页面

  2. 点击创建集群

  3. 为集群指定名称。该名称在项目和地区中必须是唯一的。

  4. 对于位置类型,请选择地区,然后选择您计划在其中使用 Cloud TPU 资源的所需地区。在本教程中,选择 us-central1-b 可用区。

  5. 务必将主版本设置为 1.13.4-gke.5 或更高版本,以支持 Cloud TPU。

  6. 在导航窗格的要配置的节点池下,点击安全

  7. 选择允许所有 Cloud API 的完整访问权限。这可确保集群中的所有节点都能访问您的 Cloud Storage 存储分区。要实现访问,该集群和存储分区必须位于同一项目中。请注意,Kubernetes Pod 默认继承其所部署到的节点的访问权限范围。如果您想要根据每个 Pod 来限制访问权限,请参阅 GKE 指南通过服务帐号进行身份验证

  8. 在导航窗格的集群下,点击网络

  9. 选择启用 VPC 原生流量路由(使用别名 IP),如果当前项目不存在 VPC 网络,您需要创建一个 VPC 网络

  10. 在导航窗格的集群下,点击功能

  11. 选择启用 Cloud TPU

  12. 根据需要配置集群的其余选项。将其余选项保留为默认值。

  13. 点击创建

  14. 连接到该集群。具体做法是,从 Console Kubernetes 集群页面选择您的集群,然后点击连接按钮。系统将显示要在 Cloud Shell 中运行以进行连接的 gcloud 命令。

gcloud

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

  1. 安装 gcloud 组件,您需要这些组件来运行支持 Cloud TPU 的 GKE:

    $ gcloud components install kubectl 
  2. 使用您的 Google Cloud 项目 ID 来配置 gcloud

    $ gcloud config set project project-name
    

    project-name 替换为您的 Google Cloud 项目的名称。

    当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 GCP API 调用。

  3. 对您计划使用 Cloud TPU 资源的可用区配置 gcloud。在本教程中,请使用 us-central1-b 可用区:

    $ gcloud config set compute/zone us-central1-b
    
  4. 使用 gcloud container clusters create 命令在 GKE 上创建一个支持 Cloud TPU 的集群。在以下命令中,将 cluster-name 替换为您选择的集群名称:

    $ gcloud container clusters create cluster-name \
      --cluster-version=1.16 \
      --scopes=cloud-platform \
      --enable-ip-alias \
      --enable-tpu
    

    命令标志说明

    cluster-version
    表示该集群将使用最新的 Kubernetes 1.16 版本。您必须使用 1.13.4-gke.5 或更高版本。
    scopes
    确保集群中的所有节点都能访问您的 Cloud Storage 存储分区。要实现访问,该集群和存储分区必须位于同一项目中。请注意,Kubernetes Pod 默认继承其所部署到的节点的访问权限范围。因此,scopes=cloud-platform 向在该集群中运行的所有 Kubernetes Pod 授予 cloud-platform 范围。如果您想要根据每个 Pod 来限制访问权限,请参阅 GKE 指南通过服务帐号进行身份验证
    enable-ip-alias
    表示集群使用别名 IP 范围。如需在 GKE 上使用 Cloud TPU,这是必需的设置。
    enable-tpu
    表示集群必须支持 Cloud TPU。
    tpu-ipv4-cidr(可选,未在上文中指定)
    表示用于 Cloud TPU 的 CIDR 范围。以 IP/20 的形式指定 IP_RANGE,例如 10.100.0.0/20。如果未指定此标志,则自动分配大小为 /20 的 CIDR 范围。

    创建集群后,您应该会看到类似于以下内容的消息:

    NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
    cluster-resnet  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING
    

查看操作

启用 Cloud TPU 支持会启动更新操作。对于 区域级集群,此操作大约需要 5 分钟;对于 地区级集群,此操作大约需要 15 分钟,具体取决于集群所在的地区。

如需列出集群中所有正在运行和已完成的操作,请运行以下命令:

   $ gcloud container operations list
   

如需详细了解特定操作,请运行以下命令:

   $ gcloud container operations describe operation-id
   

operation-id 替换为特定操作的 ID。

GKE/ResNet 作业规范

本教程使用的作业规范(如下所示)申请了一个使用 TensorFlow 2.3 的抢占式 Cloud TPU v2 设备。它还会启动 TensorBoard 过程

Cloud TPU 节点的生命周期取决于请求这些节点的 Kubernetes Pod。安排 Kubernetes Pod 时会按需创建 Cloud TPU,删除 Pod 时会回收此 Cloud TPU。

apiVersion: batch/v1
kind: Job
metadata:
  name: resnet-tpu
spec:
  template:
    metadata:
      annotations:
        # The Cloud TPUs that will be created for this Job will support
        # TensorFlow 2.3. This version MUST match the
        # TensorFlow version that your model is built on.
        tf-version.cloud-tpus.google.com: "2.3"
    spec:
      restartPolicy: Never
      containers:
      - name: resnet-tpu
        # The official TensorFlow 2.3.0 image.
        # https://hub.docker.com/r/tensorflow/tensorflow
        image: tensorflow/tensorflow:2.3.0
        command:
        - bash
        - -c
        - |
          pip install tf-models-official==2.3.0
          python3 -m official.vision.image_classification.resnet.resnet_ctl_imagenet_main \
            --tpu=$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS) \
            --distribution_strategy=tpu \
            --steps_per_loop=500 \
            --log_steps=500 \
            --use_synthetic_data=true \
            --dtype=fp32 \
            --enable_tensorboard=true \
            --train_epochs=90 \
            --epochs_between_evals=1 \
            --batch_size=1024 \
            --model_dir=gs://bucket-name/resnet
        resources:
          limits:
            # Request a single Preemptible v2-8 Cloud TPU device to train the
            # model. A single v2-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/preemptible-v2: 8
      - name: tensorboard
        image: tensorflow/tensorflow:2.2.0
        command:
        - bash
        - -c
        - |
          pip install tensorboard-plugin-profile==2.3.0 cloud-tpu-client
          tensorboard --logdir=gs://bucket-name/resnet --port=6006
        ports:
        - containerPort: 6006

创建作业

请按照以下步骤在 GKE 集群中创建作业:

  1. 使用文本编辑器创建作业规范 example-job.yaml,并在上述作业规范中进行复制/粘贴。请务必将 --model_dir 参数和 tensorboard 命令中的 bucket-name 变量替换为您的存储分区名称。

  2. 运行该作业:

    $ kubectl create -f example-job.yaml
    

    此命令会创建自动安排 Pod 的作业,并应返回以下输出。

    job "resnet-tpu" created
  3. 验证 Kubernetes Pod 是否已安排且 Cloud TPU 节点是否已预配。请求 Cloud TPU 节点的 Kubernetes Pod 可能要等到 5 分钟后才会运行。在 Kubernetes Pod 安排好之前,您将会看到类似如下所示的输出。

    $ kubectl get pods -w
    
    初始输出应如下显示:
    
    NAME               READY     STATUS    RESTARTS   AGE
    resnet-tpu-cmvlf   0/1       Pending   0          1m
    

    5 分钟后,您应该会看到如下内容:

    NAME               READY     STATUS    RESTARTS   AGE
    resnet-tpu-cmvlf   1/1       Running   0          6m
    

  4. STATUS 显示 running 时,表示已开始训练。输入 CTRL-C 以退出 kubectl get pods -w 进程。

查看 Cloud TPU 和 TensorBoard 容器状态与日志

请按照以下步骤验证状态并查看 Kubernetes Pod 所使用的 Cloud TPU 和 TensorBoard 实例的日志。

  1. 按照以下步骤创建部署和服务规范:

    转到 GKE 页面

  2. 在左侧导航栏上,点击工作负载

  3. 选择作业。系统会将您转到一个包含托管 Pod 标题的页面。

  4. 托管 Pod 下,选择您的 Kubernetes Pod。系统会将您转到包含容器标题的页面。

  5. 容器下会显示容器列表。此列表包括所有 Cloud TPU 和 TensorBoard 实例。针对每个容器,系统会显示以下信息:

    1. 运行状态
    2. 指向容器日志的链接

使用 TensorBoard 来直观显示指标和分析性能

TensorBoard 是一套专用于直观地呈现 TensorFlow 指标的工具。TensorBoard 可帮助确定处理过程中的瓶颈,并提出改善性能的建议。

TensorFlow Profiler 是一款 TensorBoard 插件,用于捕获单个 Cloud TPU 上的分析数据并在 TensorBoard 上直观呈现。如需详细了解可视化工具以及您可以捕获哪些信息,请参阅 TPU 工具文档。

请按照以下步骤在 GKE 集群中运行 TensorBoard:

  1. 按照查看 TensorBoard 状态的步骤来验证 TensorBoard 实例是否在容器中运行。

  2. 执行 Port-forward 转到 TensorBoard Kubernetes Pod:

    $ kubectl port-forward pod/resnet-tpu-pod-id 6006
    

    其中 pod-id 是控制台上显示的 GKE Pod 名称的最后一组数字:Kubernetes Engine > 工作负载 > 托管 Pod。例如:resnet-tpu-wxskc

  3. 在 Cloud Shell 右上角的边栏上,点击网页预览按钮并打开端口 6006 以查看 TensorBoard 输出。TensorBoard 界面在浏览器中将以标签页的形式显示。

  4. 从 TensorBoard 页面右上角的下拉菜单中选择 PROFILE

  5. 点击性能剖析文件页面上的捕获性能剖析文件按钮。

  6. 在弹出式菜单中,选择 TPU 名称地址类型,然后输入 TPU 名称。TPU 名称显示在 Cloud Console 中的 Compute Engine > TPU 页面,并采用以下格式:

    gke-cluster-name-cluster-id-tpu-tpu-id
    
    例如:
    gke-demo-cluster-25cee208-tpu-4b90f4c5

  7. 如果您已准备好开始分析,请在弹出菜单上选择捕获按钮,然后等待几秒钟让分析完成。

  8. 刷新浏览器,在 TensorBoard 的 PROFILE 标签页下查看跟踪记录数据。

如需详细了解如何捕获和解释性能剖析文件,请参阅 TensorFlow Profiler 指南

清理

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

控制台

删除 GKE 集群:

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

    转到 GKE 页面

  2. 选择要删除的集群旁边的复选框。

  3. 点击删除

完成数据检查后,删除在本教程期间创建的 Cloud Storage 存储分区:

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

    转到 Cloud Storage 页面

  2. 选择要删除的存储分区旁边的复选框。

  3. 点击删除

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

gcloud

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

    $ gcloud container clusters delete cluster-name --zone=us-central1-b --project=project-name
    

    此命令会删除集群、容器和 Cloud TPU。

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

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

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

后续步骤

  • 使用以下作业规范之一运行更多模型和数据集检索作业: