使用 TensorFlow 在 Cloud TPU 虚拟机上运行计算
本快速入门将向您介绍如何创建 Cloud TPU、安装 TensorFlow 以及在 Cloud TPU 上运行计算。如需更深入地了解如何在 Cloud TPU 上训练模型,请参阅 Cloud TPU 教程之一。
准备工作
在按照本快速入门进行操作之前,您必须先创建 Google Cloud 帐号、安装 Google Cloud CLI 并配置 gcloud
命令。如需了解详情,请参阅设置账号和 Cloud TPU 项目
使用 gcloud
创建 Cloud TPU 虚拟机
使用 gcloud
命令启动 Compute Engine Cloud TPU。如需详细了解 gcloud
命令,请参阅 gcloud
参考文档。
启动 TPU 虚拟机时,您可以使用命令中显示的默认 TPU 软件版本,也可以参阅 Cloud TPU 软件版本来指定其他受支持的 TPU 软件版本。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=europe-west4-a
--accelerator-type=v3-8
--version=tpu-vm-tf-2.16.1-pjrt
命令标志说明
连接到 Cloud TPU 虚拟机
您必须使用 SSH 明确连接到您的 TPU 虚拟机。如果未自动连接,请使用以下命令。
$ gcloud compute tpus tpu-vm ssh tpu-name
--zone europe-west4-a
使用 TensorFlow 运行示例
连接到 TPU 虚拟机后,请设置以下环境变量。
(vm)$ export TPU_NAME=local
创建 TPU 时,如果将 --version
参数设置为以 -pjrt
结尾的版本,请设置以下环境变量以启用 PJRT 运行时:
(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
在当前目录中创建一个名为 tpu-test.py
的文件,并将以下脚本复制粘贴到其中。
import tensorflow as tf print("Tensorflow version " + tf.__version__) @tf.function def add_fn(x,y): z = x + y return z cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(cluster_resolver) tf.tpu.experimental.initialize_tpu_system(cluster_resolver) strategy = tf.distribute.TPUStrategy(cluster_resolver) 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) }
清理
为避免因本页中使用的资源导致您的 Google Cloud 账号产生费用,请按照以下步骤操作。
断开与 Compute Engine 实例的连接(如果您尚未这样做):
(vm)$ exit
您的提示符现在应为
username@projectname
,表明您位于 Cloud Shell 中。删除您的 Cloud TPU。
$ gcloud compute tpus tpu-vm delete tpu-name
--zone=europe-west4-a通过运行
gcloud compute tpus tpu-vm list
验证资源是否已删除。删除操作可能需要几分钟时间才能完成。$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
后续步骤
如需详细了解 Cloud TPU,请参阅: