使用 PyTorch 在 Cloud TPU 上训练 Resnet50

本教程介绍如何使用 PyTorch 在 Cloud TPU 设备上训练 ResNet-50 模型。您可以将同一模式应用于使用 PyTorch 和 ImageNet 数据集的其他针对 TPU 进行了优化的图片分类模型。

本教程中的模型基于用于图片识别的深度残差学习,率先引入了残差网络 (ResNet) 架构。本教程使用 50 层变体 ResNet-50,演示如何使用 PyTorch/XLA 训练模型。

目标

  • 准备数据集。
  • 运行训练作业。
  • 验证输出结果。

费用

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

  • Compute Engine
  • Cloud TPU

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

准备工作

在开始学习本教程之前,请检查您的 Google Cloud 项目是否已正确设置。

  1. 登录您的 Google 帐号。

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

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

    转到项目选择器页面

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

  4. 本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源,以免产生不必要的费用。

设置 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. 设置 PyTorch 版本

    export VERSION=1.6
    
  5. 启动本教程所需的 Compute Engine 资源。

    gcloud compute instances create resnet50-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-highmem-96  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=300GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    

    如果您想要使用真实数据训练 Resnet50,请尽可能选择 CPU 数量最多的机器类型。Resnet50 的输入通常高度受限制,因此除非有很多工作器馈入数据和足够的 RAM 来维持大量工作器线程,否则训练速度可能会很慢。为获得最佳结果,请选择 n1-highmem-96 机器类型。

    如果您想要下载 ImageNet,请指定至少为 300GB 的磁盘大小。如果您想要仅使用虚构数据,您可以指定 20GB 的默认磁盘大小。本教程建议同时使用这两个数据集。

启动 Cloud TPU 资源

  1. 在 Compute Engine 虚拟机中,使用以下命令启动 Cloud TPU 资源:

    (vm) $ gcloud compute tpus create resnet50-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-${VERSION?}  \
    --accelerator-type=v3-8
    
  2. 确定 Cloud TPU 资源的 IP 地址。

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

创建并配置 PyTorch 环境

  1. 连接到新的 Compute Engine 实例。

    gcloud compute ssh resnet50-tutorial --zone=us-central1-a
    

从现在开始,前缀 (vm)$ 表示您应该在 Compute Engine 虚拟机实例上运行该命令。

  1. 启动 conda 环境。

    (vm) $ export VERSION=1.6
    (vm) $ conda activate torch-xla-${VERSION?}
    
  2. 为 Cloud TPU 资源配置环境变量。

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

使用虚构数据集训练

我们建议您在初始运行时使用虚构数据集,而不是真实的 ImageNet 数据集,这是因为 fake_data 会自动安装在您的虚拟机中,并且处理时所需的时间和资源较少。

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=2 --batch_size=128 --log_steps=20

使用真实数据集训练

如果使用 --fake_data 标志未出现任何问题,您可以尝试使用真实数据训练,例如 ImageNet

通常,test_train_mp_imagenet.py 使用 torchvision.datasets.ImageFolder,因此您可以使用任何结构正确的数据集。请参阅 ImageFolder 文档。

如果您将数据集存储在 ~/imagenet,则建议针对使用真实数据进行一些命令行修改:

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=8 --batch_size=128 --log_steps=200

申请更多 CPU 配额

请提前几天安排申请更多资源,以确保我们有足够的时间来处理您的申请。

  1. 转到配额页面。

    转到“配额”页面

  2. 服务菜单中选择 Cloud TPU API
  3. 选择要在其中使用 CPU 的区域或地区。
  4. 指标菜单中,选择“无”,然后在搜索框中输入 CPU。
  5. 选择 CPU
  6. 在列表中,选择 Compute Engine API - CPU,然后点击页面顶部的修改配额
  7. 输入您申请的配额量和说明(必填),然后点击完成。申请会发送给相应的服务提供商以获取批准。

清理

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

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm)$ exit
    

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

  2. 在 Cloud Shell 中,使用您在设置 Compute Engine 虚拟机和 Cloud TPU 时使用的 --zone--name 标志运行 ctpu delete。这会同时删除您的虚拟机和 Cloud TPU。

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

    $ ctpu status --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:             --
    

后续步骤

试用 PyTorch Colab: