使用 PyTorch 在 Cloud TPU 上训练 Resnet50

---

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

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

警告：本教程使用第三方数据集。Google 不对此数据集的有效性或任何其他方面提供任何代理、担保或其他保证。

目标

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

费用

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

• Compute Engine
• Cloud TPU

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

准备工作

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

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

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

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

   转到“项目选择器”

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

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

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

   转到“项目选择器”

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

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

设置 Compute Engine 实例

1. 打开一个 Cloud Shell 窗口。

   打开 Cloud Shell

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

   export PROJECT_ID=project-id
   

3. 配置 Google Cloud CLI，以使用要在其中创建 Cloud TPU 的项目。

   gcloud config set project ${PROJECT_ID}
   

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

4. 启动本教程所需的 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
   

   注意：运行此命令时，您将看到一条警告：

   Disk size: '300 GB' is larger than image size
   

   您可以放心地忽略该警告消息。

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

   注意：检查您的 CPU 配额，确保有足够的可用 CPU。如需增加配额，请参阅 CPU 配额。

   如果您想要下载 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-1.13 \
   --accelerator-type=v3-8
   

2. 确定 Cloud TPU 资源的 IP 地址。

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

   重要提示：IP 地址位于 NETWORK_ENDPOINTS 列下。 在创建和配置 PyTorch 环境时，您将需要该 IP 地址（不包括端口号）。

创建并配置 PyTorch 环境

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

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

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

1. 启动 conda 环境。

   (vm) $ conda activate torch-xla-1.13
   

2. 为 Cloud TPU 资源配置环境变量。

   注意：TPU_IP_ADDRESS 变量必须与您在启动 Cloud TPU 资源时确定的 Cloud TPU 的 IP 地址相等（不包括端口号）。它是 Google Cloud Console 上 Compute Engine > TPU 下显示的内部 IP 地址。

   (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-1.13/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-1.13/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 帐号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

1. 断开与 Compute Engine 实例的连接（如果您尚未这样做）：

   (vm)$ exit
   

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

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

   $ ctpu delete --project=${PROJECT_ID} \
     --name=resnet50-tutorial \
     --zone=us-central1-a
   

   重要提示：如果您在运行 ctpu up 时设置了 TPU 资源名称，则必须在运行 ctpu delete 时使用 --name 标志指定该名称才能关停 TPU 资源。

3. 运行 ctpu status 以确保未分配任何实例，避免产生不必要的 TPU 使用费。删除操作可能需要几分钟时间才能完成。 如下所示的响应表明不再有已分配的实例：

   $ ctpu status --project=${PROJECT_ID} \
     --zone=us-central1-a
   

   2018/04/28 16:16:23 WARNING: Setting zone to "us-central1-a"
   No instances currently exist.
           Compute Engine VM:     --
           Cloud TPU:             --
   

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

后续步骤

试用 PyTorch Colab：

• 在 Cloud TPU 上开始使用 PyTorch
• 在 TPU 上训练 MNIST
• 使用 Cifar10 数据集在 TPU 上训练 ResNet18
• 使用预训练的 ResNet50 模型进行推理
• 快速神经风格转移
• 在 Fashion MNIST 上使用多核心训练 AlexNet
• 在 Fashion MNIST 上使用单核心训练 AlexNet