在 TPU Pod 切片上运行 TensorFlow 代码

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

  1. 使用 TensorFlow 软件创建 TPU Pod 切片
  2. 使用 SSH 连接到 TPU 虚拟机
  3. 创建并运行简单的脚本

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

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

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.17.0-pod-pjrt

命令标志说明

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

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

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

创建并运行简单的计算脚本

  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 的每个 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 europe-west4-a

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

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