在 TPU Pod 切片上运行 TensorFlow 代码
本文档介绍如何在 TPU Pod 上使用 TensorFlow 执行简单的计算。您将执行以下步骤:
- 使用 TensorFlow 软件创建 TPU Pod 切片
- 使用 SSH 连接到 TPU 虚拟机
- 创建并运行简单的脚本
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
创建并运行简单的计算脚本
设置以下环境变量。
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
在当前目录中创建一个名为
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 虚拟机的操作后,请按照以下步骤清理资源。
断开与 Compute Engine 的连接:
(vm)$ exit
删除您的 Cloud TPU。
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
通过运行以下命令来验证资源已删除。确保您的 TPU 不再列出。删除操作可能需要几分钟时间才能完成。
$ gcloud compute tpus tpu-vm list \ --zone europe-west4-a