概览
要使用 Cloud TPU,您需要选择软件版本以匹配在本地虚拟机上运行的软件。您现在可以在运行中的 Cloud TPU 上切换此软件版本,而无需删除和重新创建 TPU。您还还可以将 Cloud TPU 配置为使用特定每夜版的软件框架。我们仍然建议您选择这些框架的受支持版本。
用量
切换版本的推荐方法是使用 cloud-tpu-client Python 库。
TensorFlow 的示例用法。
import tensorflow as tf
from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version(tf.__version__, restart_type='ifNeeded')
这会将 Cloud TPU 配置为与本地虚拟机上运行的 TensorFlow 版本完全匹配,包括官方版本以及过时的每夜版。
该库直接与 Cloud TPU 通信,因此该代码需要在同一网络中的虚拟机中运行。建议您在模型的其余部分的代码中运行此代码。
其他软件选项
TensorFlow 包含一个 tf.__version__
字符串,这是配置正确版本的最简单方式。其他软件选项包括:
- PyTorch -
pytorch-0.1
、pytorch-0.5-dev20200320
、pytorch-nightly
- Jax -
tpu_driver
、tpu_driver0.1-dev20200320
、tpu_driver_nightly
例如,将 TPU 配置为使用 PyTorch 的最新每夜版运行。
from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()