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。

import os
import pprint
import tensorflow as tf
 
if 'COLAB_TPU_ADDR' not in os.environ:
    print('ERROR: Not connected to a TPU runtime')
else:
    tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
    print ('TPU address is', tpu_address)
 
    with tf.Session(tpu_address) as session:
      devices = session.list_devices()
 
    print('TPU devices:')
    pprint.pprint(devices)

输出信息：

TPU address is grpc://10.49.237.2:8470
TPU devices:
[_DeviceAttributes(/job:tpu_worker/.../device:CPU:0, CPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:XLA_CPU:0, XLA_CPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:0, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:1, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:2, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:3, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:4, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:5, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:6, TPU, ...),
 _DeviceAttributes(/job:tpu_worker/.../device:TPU:7, TPU, ...),
 _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了。

> Source: TPUs for Developers

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