在 TPU Pod 切片上运行 TensorFlow 代码

设置运行 TensorFlow 的 TPU 虚拟机 Pod 并运行计算

1. 按照以下步骤设置运行 TensorFlow 的 Pod 并运行计算。

   注意：启动 TPU 虚拟机时，您可以使用命令中显示的 TPU 软件版本，也可以参阅 Cloud TPU 软件版本来指定其他受支持的 TPU 软件版本。

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

   命令标志说明

   zone

   您计划在其中创建 Cloud TPU 的可用区。

   accelerator-type

   要创建的 Cloud TPU 的类型。

   version

   Cloud TPU 软件版本。默认选项为最新的 TensorFlow 软件版本。

2. 连接到 Cloud TPU 虚拟机

   通过 SSH 连接到 Cloud TPU Pod 中的 TPU 虚拟机：

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

3. 设置以下环境变量。

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

4. 在当前目录中创建一个名为 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)

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

(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