在 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.13.0-pod

命令标志说明

zone
您计划创建 Cloud TPU 的可用区。
accelerator-type
要创建的 Cloud TPU 的类型。
version
Cloud TPU 软件 版本。默认值为最新的 TensorFlow 软件版本。

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