本教程介绍如何使用 PyTorch 在 Cloud TPU 设备上训练 ResNet-50 模型。您可以将同一模式应用于使用 PyTorch 和 ImageNet 数据集的其他针对 TPU 进行了优化的图片分类模型。
本教程中的模型基于用于图片识别的深度残差学习,率先引入了残差网络 (ResNet) 架构。本教程使用 50 层变体 ResNet-50,演示如何使用 PyTorch/XLA 训练模型。
目标
- 准备数据集。
- 运行训练作业。
- 验证输出结果。
费用
本教程使用 Google Cloud 的以下收费组件:
- Compute Engine
- Cloud TPU
您可使用价格计算器根据您的预计使用情况来估算费用。
准备工作
在开始学习本教程之前,请检查您的 Google Cloud 项目是否已正确设置。
- 登录您的 Google Cloud 帐号。如果您是 Google Cloud 新手,请创建一个帐号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
-
在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目。
-
确保您的 Cloud 项目已启用结算功能。了解如何检查项目是否已启用结算功能。
-
在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目。
-
确保您的 Cloud 项目已启用结算功能。了解如何检查项目是否已启用结算功能。
本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源,以免产生不必要的费用。
设置 Compute Engine 实例
打开一个 Cloud Shell 窗口。
为项目 ID 创建一个变量。
export PROJECT_ID=project-id
配置 Google Cloud CLI,以使用要在其中创建 Cloud TPU 的项目。
gcloud config set project ${PROJECT_ID}
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示
Authorize Cloud Shell
页面。点击页面底部的Authorize
,以允许gcloud
使用您的凭据进行 Google Cloud API 调用。启动本教程所需的 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 资源
在 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
确定 Cloud TPU 资源的 IP 地址。
(vm) $ gcloud compute tpus list --zone=us-central1-a
创建并配置 PyTorch 环境
连接到新的 Compute Engine 实例。
gcloud compute ssh resnet50-tutorial --zone=us-central1-a
从现在开始,前缀 (vm)$
表示您应该在 Compute Engine 虚拟机实例上运行该命令。
启动
conda
环境。(vm) $ conda activate torch-xla-1.13
为 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-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 配额
请提前几天安排申请更多资源,以确保我们有足够的时间来处理您的申请。
- 转到配额页面。
- 从服务菜单中选择 Cloud TPU API。
- 选择要在其中使用 CPU 的区域或地区。
- 在指标菜单中,选择“无”,然后在搜索框中输入 CPU。
- 选择 CPU。
- 在列表中,选择 Compute Engine API - CPU,然后点击页面顶部的修改配额。
- 输入您申请的配额量和说明(必填),然后点击完成。申请会发送给相应的服务提供商以获取批准。
清理
为避免因本教程中使用的资源导致您的 Google Cloud 帐号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。
断开与 Compute Engine 实例的连接(如果您尚未这样做):
(vm)$ exit
您的提示符现在应为
user@projectname
,表明您位于 Cloud Shell 中。在 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 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: --
后续步骤
试用 PyTorch Colab:
- 在 Cloud TPU 上开始使用 PyTorch
- 在 TPU 上训练 MNIST
- 使用 Cifar10 数据集在 TPU 上训练 ResNet18
- 使用预训练的 ResNet50 模型进行推理
- 快速神经风格转移
- 在 Fashion MNIST 上使用多核心训练 AlexNet
- 在 Fashion MNIST 上使用单核心训练 AlexNet