使用 TensorFlow 在 Cloud TPU 虚拟机上运行计算

本快速入门介绍了如何创建 Cloud TPU，以及如何使用 TensorFlow 在 Cloud TPU 上运行计算。如需更深入地了解如何在 Cloud TPU 上训练模型，请参阅 Cloud TPU 教程之一。

准备工作

在按照本快速入门操作之前，您必须创建一个 Google Cloud 账号，安装 Google Cloud CLI，并配置 gcloud 命令。如需了解详情，请参阅设置账号和 Cloud TPU 项目。

使用 gcloud 创建 Cloud TPU 虚拟机

使用 gcloud 命令创建 Cloud TPU。

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-8 \
  --version=tpu-vm-tf-2.18.0-pjrt \
  --project=your-gcp-project-name

命令标志说明

tpu-name
要创建的 Cloud TPU 的名称。
zone
拟在其中创建 Cloud TPU 的可用区。
accelerator-type
加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
version
TPU 运行时版本。 本快速入门中使用的版本已预安装 TensorFlow。
project
您要在其中创建 Cloud TPU 的 Google Cloud CLI 项目的名称。

如需详细了解 gcloud 命令，请参阅 gcloud 参考文档。

连接到 Cloud TPU 虚拟机

使用 SSH 连接到 TPU 虚拟机：

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a \
  --project=your-gcp-project-name

验证 TensorFlow 是否可以访问 TPU

  1. 在当前目录中创建一个名为 tpu-count.py 的文件，并将以下脚本复制粘贴到其中。

    import tensorflow as tf

print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")

  2. 运行脚本：

    (vm)$ python3 tpu-count.py

    脚本的输出会显示 TPU 虚拟机可用的 TPU 核心数：

    TensorFlow can access 8 TPU cores

使用 TensorFlow 运行基本计算

连接到 TPU 虚拟机后，请设置以下环境变量。

  (vm)$ export TPU_NAME=local

创建 TPU 时，如果您将 --version 参数设置为以 -pjrt 结尾的版本，请设置以下环境变量以启用 PJRT 运行时： 

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

在当前目录中创建一个名为 tpu-test.py 的文件，并将以下脚本复制粘贴到其中。

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

@tf.function
def add_fn(x,y):
  z = x + y
  return z

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.TPUStrategy(cluster_resolver)

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

使用以下命令运行此脚本：

(vm)$ python3 tpu-test.py

此脚本对 TPU 的每个 TensorCore 执行计算。输出将如下所示：

PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32)
}

清理

为避免因本页中使用的资源导致您的 Google Cloud 账号产生费用，请按照以下步骤操作。

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

    (vm)$ exit

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

  2. 删除您的 Cloud TPU。

    $ gcloud compute tpus tpu-vm delete tpu-name \
--zone=europe-west4-a

  3. 通过运行 gcloud compute tpus tpu-vm list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。

    $ gcloud compute tpus tpu-vm list --zone=europe-west4-a

后续步骤

如需详细了解 Cloud TPU，请参阅：