Cloud TPU 虚拟机 PyTorch 快速入门

概览:本快速入门简要介绍如何搭配使用 PyTorch 和 Cloud TPU。

准备工作

在开始某个 Cloud TPU 虚拟机快速入门之前,请确保已阅读 Cloud TPU 用户指南。本部分介绍如何设置 Google Cloud 项目以及如何创建 Cloud TPU 虚拟机。

登录您的 Google 帐号。 如果您还没有 Google 帐号,请注册新帐号。在 Google Cloud Console 中,从项目选择器页面选择或创建 Cloud 项目。确保您的项目已启用结算功能。

启用 Cloud TPU API

Cloud Shell 中使用以下 gcloud 命令启用 Cloud TPU API。(您也可以从 Google Cloud Console 启用。)

$ gcloud services enable tpu.googleapis.com

配置 gcloud 命令

运行以下命令配置 gcloud,以使用 GCP 项目并安装 TPU 虚拟机预览版所需的组件。

$ gcloud config set account your-email-account
$ gcloud config set project project-id

使用 gcloud 创建 Cloud TPU 虚拟机

$ gcloud alpha compute tpus tpu-vm create tpu-name \
  --zone=us-central1-b \
  --accelerator-type=v3-8 \
  --version=v2-alpha

必需的字段

zone
拟在其中创建 Cloud TPU 的区域
accelerator-type
要创建的 Cloud TPU 的类型
version
Cloud TPU 运行时版本。

连接到 Cloud TPU 虚拟机

$ gcloud alpha compute tpus tpu-vm ssh tpu-name --zone us-central1-b --project project-id

必需的字段

tpu_name
要连接的 TPU 虚拟机的名称。
zone
您要在其中创建 Cloud TPU 的可用区。目前支持 europe-west4-aus-central1-a
project-id
您的 GCP 项目 ID。

设置 XRT TPU 设备配置

$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

执行简单的计算:

启动 Python 3 解释器:

(vm)$ python3

命令的输出为:

Python 3.6.9 (default, Oct  8 2020, 12:12:24)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.

导入 PyTorch 软件包

>>> import torch
>>> import torch_xla.core.xla_model as xm

执行计算

>>> dev = xm.xla_device()
>>> t1 = torch.randn(3,3,device=dev)
>>> t2 = torch.randn(3,3,device=dev)
>>> print(t1 + t2)

命令的输出为:

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

退出 Python 解释器:

>>> exit()

清除数据

为避免系统因本快速入门中使用的资源向您的 Google Cloud 帐号收取费用,请按照以下步骤操作。

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

    (vm)$ exit

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

  2. 删除您的 Cloud TPU。

    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone=us-central1-b
  3. 通过运行 gcloud alpha compute tpus tpu-vm list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

     $ gcloud alpha compute tpus tpu-vm list --zone=us-central1-b

     NAME             STATUS

后续步骤

详细了解 Cloud TPU 虚拟机: