在 TPU Pod 切片上运行 TensorFlow 代码

本文档介绍了如何在 TPU Pod 上使用 TensorFlow 执行计算。您将执行以下步骤：

  1. 使用 TensorFlow 软件创建 TPU Pod 切片
  2. 使用 SSH 连接到 TPU 虚拟机
  3. 创建和运行示例脚本

TPU 虚拟机依赖于服务账号 获取调用 Cloud TPU API 的权限。默认情况下，TPU 虚拟机将使用默认的 Compute Engine 服务账号，其中包含所有所需的 Cloud TPU 权限。如果您使用的是自己的 服务账号，您需要添加 TPU Viewer 角色。如需详细了解 Google Cloud 角色，请参阅了解角色。 在以下情况下，您可以使用 --service-account 标志指定自己的服务账号： 创建 TPU 虚拟机的过程。

设置环境

  1. 在 Cloud Shell 中，运行以下命令以确保 运行当前版本的 gcloud：

    $ gcloud components update

    如果您需要安装 gcloud，请使用以下命令：

    $ sudo apt install -y google-cloud-sdk

  2. 创建一些环境变量：

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

使用 TensorFlow 运行时创建 v3-32 TPU Pod 切片

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

命令标志说明

zone
拟在其中创建 Cloud TPU 的区域。
accelerator-type
加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
version
Cloud TPU 软件版本。

使用 SSH 连接到您的 Cloud TPU 虚拟机

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

创建并运行示例脚本

  1. 设置以下环境变量。

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

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

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

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

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

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

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

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

    (vm)$ python3 tpu-test.py

    此脚本会对 TPU Pod 切片的每个 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),
  8: tf.Tensor(2.0, shape=(), dtype=float32),
  9: tf.Tensor(2.0, shape=(), dtype=float32),
  10: tf.Tensor(2.0, shape=(), dtype=float32),
  11: tf.Tensor(2.0, shape=(), dtype=float32),
  12: tf.Tensor(2.0, shape=(), dtype=float32),
  13: tf.Tensor(2.0, shape=(), dtype=float32),
  14: tf.Tensor(2.0, shape=(), dtype=float32),
  15: tf.Tensor(2.0, shape=(), dtype=float32),
  16: tf.Tensor(2.0, shape=(), dtype=float32),
  17: tf.Tensor(2.0, shape=(), dtype=float32),
  18: tf.Tensor(2.0, shape=(), dtype=float32),
  19: tf.Tensor(2.0, shape=(), dtype=float32),
  20: tf.Tensor(2.0, shape=(), dtype=float32),
  21: tf.Tensor(2.0, shape=(), dtype=float32),
  22: tf.Tensor(2.0, shape=(), dtype=float32),
  23: tf.Tensor(2.0, shape=(), dtype=float32),
  24: tf.Tensor(2.0, shape=(), dtype=float32),
  25: tf.Tensor(2.0, shape=(), dtype=float32),
  26: tf.Tensor(2.0, shape=(), dtype=float32),
  27: tf.Tensor(2.0, shape=(), dtype=float32),
  28: tf.Tensor(2.0, shape=(), dtype=float32),
  29: tf.Tensor(2.0, shape=(), dtype=float32),
  30: tf.Tensor(2.0, shape=(), dtype=float32),
  31: tf.Tensor(2.0, shape=(), dtype=float32)
}

清理

完成 TPU 虚拟机的操作后，请按照以下步骤清理资源。

  1. 与 Compute Engine 断开连接：

    (vm)$ exit

  2. 删除您的 Cloud TPU。

    $ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --zone=${ZONE}

  3. 通过运行以下命令来验证资源已删除。确保您的 TPU 不再列出。删除操作可能需要几分钟时间才能完成。

    $ gcloud compute tpus tpu-vm list \
  --zone=${ZONE}