快速入门：使用 TensorFlow 在 Cloud TPU 虚拟机上运行计算

使用 TensorFlow 在 Cloud TPU 虚拟机上运行计算

本快速入门介绍如何创建 Cloud TPU、安装 TensorFlow 以及在 Cloud TPU 上运行简单的计算。如需获得更深入的教程来了解如何在 Cloud TPU 上训练模型，请参阅 Cloud TPU 教程之一。

准备工作

按照本快速入门操作之前，您必须先创建一个 Google Cloud Platform 帐号，安装 Google Cloud CLI 并配置 gcloud 命令。如需了解详情，请参阅设置帐号和 Cloud TPU 项目

使用 gcloud 创建 Cloud TPU 虚拟机或节点

使用 gcloud 命令启动 Compute Engine Cloud TPU。使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需详细了解这两种虚拟机架构，请参阅系统架构。如需详细了解 gcloud 命令，请参阅 gcloud 参考文档。

TPU 虚拟机

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

命令标志说明

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

TPU 节点

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.11.0 \

命令标志说明

project
您的 GCP 项目 ID
name
要创建的 Cloud TPU 的名称。
zone
您计划在其中创建 Cloud TPU 的可用区。
disk-size
gcloud 命令所创建虚拟机的硬盘大小（以 GB 为单位）。
machine-type
要创建的 Compute Engine 虚拟机的机器类型。
tf-version
虚拟机上安装的 TensorFlow 版本gcloud。请参阅
accelerator-type
要创建的 Cloud TPU 的类型。

连接到 Cloud TPU 虚拟机

使用 TPU 虚拟机时，您必须通过 SSH 明确连接到 TPU 虚拟机。使用 TPU 节点时，您应该通过 SSHed 自动连接到 Compute Engine 虚拟机。如果未自动连接，请使用以下命令。

TPU 虚拟机

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

TPU 节点

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

使用 TensorFlow 运行简单示例

TPU 虚拟机

连接到 TPU 虚拟机后，设置以下环境变量。

(vm)$ export TPU_NAME=local

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

TPU 节点

在当前目录中创建一个名为 tpu-test.py 的文件，并将以下脚本复制粘贴到其中。

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

@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)
}

清理

为避免因本页中使用的资源导致您的 Google Cloud 帐号产生费用，请按照以下步骤操作。

  1. 断开与 Compute Engine 实例的连接（如果您尚未这样做）：

    (vm)$ exit

    您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

  2. 删除您的 Cloud TPU。

    TPU 虚拟机

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

    TPU 节点

    $ gcloud compute tpus execution-groups delete tpu-name \
--zone=europe-west4-a

  3. 通过运行 gcloud compute tpus tpu-vm list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。

    TPU 虚拟机

    $ gcloud compute tpus tpu-vm list --zone=europe-west4-a

    TPU 节点

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

后续步骤

如需详细了解 Cloud TPU，请参阅：