Cloud TPU 虚拟机 TensorFlow 快速入门

概览:本快速入门简要介绍如何使用 Cloud TPU。在本快速入门中,您将使用 Cloud TPU 来训练 ResNet

准备工作

设置 GCP 项目

登录您的 Google 帐号。 如果您还没有 Google 帐号,请注册新帐号。在 Google Cloud Console 中,从项目选择器页面选择或创建 Cloud 项目。确保您的项目已启用结算功能。在 Cloud Shell 中使用 gcloud 设置项目 ID。项目 ID 是 Cloud Console 中显示的项目的名称。

$ gcloud config set project project-id

启用 Cloud TPU API

Cloud Shell 中使用以下 gcloud 命令启用 Cloud TPU API。(您也可以从 Google Cloud Console 启用。)

$ gcloud services enable tpu.googleapis.com

配置 gcloud 命令

运行以下命令配置 gcloud,以使用 GCP 项目并安装 TPU 虚拟机预览版所需的组件。

$ gcloud config set account your-email-account
$ gcloud config set project project-id

使用 gcloud 创建 Cloud TPU 虚拟机

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

必需的字段

zone
拟在其中创建 Cloud TPU 的区域
accelerator-type
要创建的 Cloud TPU 的类型
version
Cloud TPU 运行时版本。

连接到 Cloud TPU 虚拟机

$ gcloud alpha compute tpus tpu-vm ssh tpu-name --zone europe-west4-a --project project-id

必需的字段

tpu_name
要连接的 TPU 虚拟机的名称。
zone
您要在其中创建 Cloud TPU 的可用区。目前支持 europe-west4-aus-central1-a
project-id
您的 GCP 项目 ID。

在 Cloud TPU 虚拟机上安装 TensorFlow

(vm)$ git clone https://github.com/tensorflow/models.git
(vm)$ pip3 install -r models/official/requirements.txt

系统检查

运行一个简单测试,该测试直接使用 libtpu 库将缓冲区加载到 TPU 并执行简单计算。仅单设备 TPU 支持此命令。

$ /usr/share/tpu/libtpu_client

您会看到以下输出内容:

------ Going to Query Version ------
TPU Driver Version: 0.0.0-dev
------ Going to Open a TPU Driver ------
------ Going to Query for System Information ------
------ Going to Compile a TPU program ------
------ Going to Load a TPU program ------
------ Going to Allocate a TPU Buffer ------
------ Going to Allocate a TPU Buffer ------
------ Going to Allocate a TPU Buffer ------
------ Going to Transfer To Device ------
------ Going to Transfer To Device ------
------ Going to Execute a TPU program ------
------ Going to Transfer From Device ------
------ Going to Unload a TPU program ------
------ Going to Deallocate a TPU Buffer ------
------ Going to Deallocate a TPU Buffer ------
------ Going to Deallocate a TPU Buffer ------
sum:
3 3 3 3 3 3 3 3 3 3 [...]

使用 TensorFlow 运行简单示例

此示例在 TPU 上执行简单的计算。

$ python3 /usr/share/tpu/tensorflow/simple_example.py

在单个 TPU 上训练 Resnet

$ export PYTHONPATH=/usr/share/tpu/tensorflow/resnet50_keras
$ python3 /usr/share/tpu/tensorflow/resnet50_keras/resnet50_single_tpu.py

清除数据

为避免系统因本快速入门中使用的资源向您的 Google Cloud 帐号收取费用,请按照以下步骤操作。

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm)$ exit
    

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

  2. 删除您的 Cloud TPU。

    $ gcloud alpha compute tpus tpu-vm delete tpu_name \
      --zone=europe-west4-a
    
  3. 通过运行 gcloud alpha compute tpus tpu-vm list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

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

后续步骤

详细了解 Cloud TPU 虚拟机: