本教程介绍如何使用 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 项目。
-
在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目。
本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源,以免产生不必要的费用。
创建 TPU 虚拟机
打开一个 Cloud Shell 窗口。
创建 TPU 虚拟机
gcloud compute tpus tpu-vm create your-tpu-name \ --accelerator-type=v4-8 \ --version=tpu-vm-v4-pt-2.0 \ --zone=us-central2-b \ --project=your-project
使用 SSH 连接到您的 TPU 虚拟机:
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=us-central2-b
-
(vm)$ git clone --depth=1 --branch r2.0 https://github.com/pytorch/xla.git
使用虚构数据运行训练脚本
(vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1
如果您能够使用虚构数据训练模型,您可以尝试使用真实数据(例如 ImageNet)进行训练。如需了解如何下载 ImageNet,请参阅下载 ImageNet。在训练脚本命令中,--datadir
标志指定要用于训练的数据集的位置。以下命令假定 ImageNet 数据集位于 ~/imagenet
。
(vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --batch_size=256 --num_epochs=1
清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。
断开与 TPU 虚拟机的连接:
(vm) $ exit
您的提示符现在应为
username@projectname
,表明您位于 Cloud Shell 中。删除您的 TPU 虚拟机。
$ gcloud compute tpus tpu-vm delete resnet50-tutorial \ --zone=us-central2-b
后续步骤
试用 PyTorch Colab:
- 在 Cloud TPU 上开始使用 PyTorch
- 在 TPU 上训练 MNIST
- 使用 Cifar10 数据集在 TPU 上训练 ResNet18
- 使用预训练的 ResNet50 模型进行推理
- 快速神经风格转移
- 在 Fashion MNIST 上使用多核心训练 AlexNet
- 在 Fashion MNIST 上使用单核心训练 AlexNet