在 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 角色，请参阅了解角色。 您可以在创建 TPU 虚拟机时使用 --service-account 标志指定自己的服务帐号。

使用 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.16.1-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