TPU 环境配置

免费 TPU：Google Colab

最方便使用 TPU 的方法，就是使用 Google 的 Colab ，不但通过浏览器访问直接可以用，而且还免费。

在 Google Colab 的 Notebook 界面中，打开界面中，打开主菜单 Runtime ，然后选择 Change runtime type，会弹出 Notebook settings 的窗口。选择里面的 Hardware accelerator 为 TPU 就可以了。

为了确认 Colab Notebook 中的确分配了 TPU 资源，我们可以运行以下测试代码。如果输出 ERROR 信息，则表示目前的 Runetime 并没有分配到 TPU；如果输出 TPU 地址及设备列表，则表示 Colab 已经分配了 TPU。

  1. import os
  2. import pprint
  3. import tensorflow as tf
  4.  
  5. if 'COLAB_TPU_ADDR' not in os.environ:
  6.     print('ERROR: Not connected to a TPU runtime')
  7. else:
  8.     tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  9.     print ('TPU address is', tpu_address)
  10.  
  11.     with tf.Session(tpu_address) as session:
  12.       devices = session.list_devices()
  13.  
  14.     print('TPU devices:')
  15.     pprint.pprint(devices)

输出信息：

  1. TPU address is grpc://10.49.237.2:8470
  2. TPU devices:
  3. [_DeviceAttributes(/job:tpu_worker/.../device:CPU:0, CPU, ...),
  4.  _DeviceAttributes(/job:tpu_worker/.../device:XLA_CPU:0, XLA_CPU, ...),
  5.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:0, TPU, ...),
  6.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:1, TPU, ...),
  7.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:2, TPU, ...),
  8.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:3, TPU, ...),
  9.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:4, TPU, ...),
  10.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:5, TPU, ...),
  11.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:6, TPU, ...),
  12.  _DeviceAttributes(/job:tpu_worker/.../device:TPU:7, TPU, ...),
  13.  _DeviceAttributes(/job:tpu_worker/.../device:TPU_SYSTEM:0, TPU_SYSTEM, ...)]

看到以上信息（一个CPU worker，8个TPU workers），既可以确认 Colab 的 TPU 环境设置正常。

Cloud TPU

在 Google Cloud 上，我们可以购买所需的 TPU 资源，用来按需进行机器学习训练。为了使用 Cloud TPU ，需要在 Google Cloud Engine 中启动 VM 并为 VM 请求 Cloud TPU 资源。请求完成后，VM 就可以直接访问分配给它专属的 Cloud TPU了。

../../_images/cloud-tpu-architecture.png

> Source: TPUs for Developers

在使用 Cloud TPU 时，为了免除繁琐的驱动安装，我们可以通过直接使用 Google Cloud 提供的 VM 操作系统镜像。