Question

这份文档说明了有效使用 Cloud TPU 时必需使用的关键 TensorFlow APIs，并强调了常规的 TensorFlow 和在 TPU 上使用区别。

这份文档针对以下用户：

• 熟悉 TensorFlow 的 Estimator 和 Dataset APIs
• 使用一个已有模型 尝试使用过 Cloud TPU
• 浏览过 TPU 模型的样例代码 [1] [2]
• 对将一个现有的 Estimator 模型移植到 Cloud TPU 上运行感兴趣

TPUEstimator

tf.contrib.tpu.TPUEstimator

创建定制化 Estimator

TPUEstimator 类和 Estimator 之间多少有些不一样。

要使一个模型可以在 CPU/GPU 或 Cloud TPU 上运行的最简单方法是在 model_fn 外定义模型的推理过程（从输入到预测）。然后继续分离 Estimator 设置和 model_fn ，都包含这个推理步骤。这种模式的一个样例是 tensorflow/models 中比较 mnist.py 和 mnist_tpu.py 的实现。

本地运行 TPUEstimator

要创建一个标准的 Estimator 你可以调用构造函数，并将它传递给 model_fn ，例如：

my_estimator = tf.estimator.Estimator(
  model_fn=my_model_fn)

tf.contrib.tpu.RunConfig

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config=tf.contrib.tpu.RunConfig()
    use_tpu=False)

这样简单的更改就能在本地运行 TPUEstimator 。大多数 TPU 模型示例都可以以下面这种命令行设置标志参数，来在本地模式下运行：

$> python mnist_tpu.py --use_tpu=false --master=''

注意： use_tpu=False 参数对于尝试 TPUEstimator API 很有用。这也就意味着它不是个完整的 TPU 兼容测试。在 TPUEstimator 中成功地本地运行一个模型并不代表它能在 TPU 上运行。

构建一个 tpu.RunConfig

虽然默认的 RunConfig 足以进行本地训练，但在实际使用并不能忽略这些设置。

一种可以切换到 Cloud TPU 的更典型 RunConfig 设置，会如下所示：

import tempfile
import subprocess

class FLAGS(object):
  use_tpu=False
  tpu_name=None
  # 为 `model_dir` 设定本地临时路径
  model_dir = tempfile.mkdtemp()
  # 在返回控制之前在 Cloud TPU 上运行的训练步数
  iterations = 50
  # 一个包含 8 个分片的 Cloud TPU
  num_shards = 8

if FLAGS.use_tpu:
    my_project_name = subprocess.check_output([
        'gcloud','config','get-value','project'])
    my_zone = subprocess.check_output([
        'gcloud','config','get-value','compute/zone'])
    cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
            tpu_names=[FLAGS.tpu_name],
            zone=my_zone,
            project=my_project)
    master = tpu_cluster_resolver.get_master()
else:
    master = ''

my_tpu_run_config = tf.contrib.tpu.RunConfig(
    master=master,
    evaluation_master=master,
    model_dir=FLAGS.model_dir,
    session_config=tf.ConfigProto(
        allow_soft_placement=True, log_device_placement=True),
    tpu_config=tf.contrib.tpu.TPUConfig(FLAGS.iterations,
                                        FLAGS.num_shards),
)

tf.contrib.tpu.RunConfig

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config = my_tpu_run_config,
    use_tpu=FLAGS.use_tpu)

通常， FLAGS 将由命令行参数设置。要从本地训练转换为 Cloud TPU 训练，你需要：

• 设置 FLAGS.use_tpu 为 True
• 设置 FLAGS.tpu_name ，以便 tf.contrib.cluster_resolver.TPUClusterResolver 可以找到它
• 设置 FLAGS.model_dir 为一个 Google Cloud Storage 容器地址（ gs:// ）。

优化器

tf.contrib.tpu.CrossShardOptimizer

CrossShardOptimizer 不兼容本地训练。因此，如果要在本地和 Cloud TPU 上运行相同的代码，请添加如下代码：

optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
if FLAGS.use_tpu:
  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

如果想在模型代码中避免使用全局 FLAGS ，一种方法就是将优化器设置为 Estimator 的参数之一，如下所示：

my_tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=my_model_fn,
    config = my_tpu_run_config,
    use_tpu=FLAGS.use_tpu,
    params={'optimizer':optimizer})

模型函数

本节详细介绍了使模型函数（ model_fn() ）能与 TPUEstimator 兼容所要做的必要更改。

静态形状

张量

$xla$XLA

摘要

将模型中所有的 tf.summary 都删除。

Tensorboard：可视化学习面板

评估标准

在一个独立的 metric_fn 中构建评估指标字典。

评估指标是训练的重要部分。Cloud TPU 完全支持这些功能，但语法略有不同。

tf.metrics

running_average, current_batch = tf.metrics.accuracy(labels, predictions)

在标准的 Estimator 中，创建这些张量对的字典，并将其作为 Estimator 的一部分返回。

my_metrics = {'accuracy': tf.metrics.accuracy(labels, predictions)}

return tf.estimator.EstimatorSpec(
  ...
  eval_metric_ops=my_metrics
)

相反，在 TPUEstimator 中，传递一个函数（返回一个度量词典）和一个参数张量列表，如下所示：

def my_metric_fn(labels, predictions):
   return {'accuracy': tf.metrics.accuracy(labels, predictions)}

return tf.contrib.tpu.TPUEstimatorSpec(
  ...
  eval_metrics=(my_metric_fn, [labels, predictions])
)

使用 TPUEstimatorSpec

TPUEstimatorSpec 不支持钩子，并且某些字段需要函数装饰器。

Estimator 的 model_fn 必须返回 EstimatorSpec 。 EstimatorSpec 是一种简单结构的命名字段，它包含模型中可能需要与 Estimator 交互的所有 tf.Tensors 。

tf.estimator.EstimatorSpec

• eval_metric_ops 必须被包装在 metrics_fn 中，这个字段会被重命名为 eval_metrics ( see above )。
• tf.train.SessionRunHook
• tf.train.Scaffold

Scaffold and Hooks 是高级用法，通常被忽略。

输入函数

因为输入函数是运行在主机上而不是 Cloud TPU 上的，所以它的运行方式没太大变化。本节主要解释了两项必要的调整。

Params 参数

标准 Estimator 的 input_fn 可以 包含一个 params 参数； TPUEstimator 的 input_fn 必须 包含一个 params 参数。这是允许估计器为输入流的每个副本设置批大小的必须参数。因此， TPUEstimator 的 input_fn 最简形式如下：

def my_input_fn(params):
  pass

params['batch-size'] 包含了批次大小

静态形状和批次大小

由 input_fn 生成的输入管道在 CPU 上运行。因此，它并不需要遵循 XLA/TPU 环境下严格的静态形状要求。只有一个要求是，从输入管道输送到 TPU 的成批数据具有静态形状，由标准 TensorFlow 形状推断算法确定。中间张量可以随意，能具有动态形状。如果形状推断失败，但已知形状，则可以使用 tf.set_shape() 强制施加正确的形状。

在下面的示例中，形状推断算法失败，但使用了 set_shape 进行了更正：

>>> x = tf.zeros(tf.constant([1,2,3])+1)
>>> x.shape

TensorShape([Dimension(None), Dimension(None), Dimension(None)])

>>> x.set_shape([2,3,4])

在许多情况下，批次大小是唯一未知的维度。

tf.data.Dataset.batch

>>> params = {'batch_size':32}
>>> ds = tf.data.Dataset.from_tensors([0, 1, 2])
>>> ds = ds.repeat().batch(params['batch-size'])
>>> ds

<BatchDataset shapes: (?, 3), types: tf.int32>

tf.contrib.data.batch_and_drop_remainder

>>> params = {'batch_size':32}
>>> ds = tf.data.Dataset.from_tensors([0, 1, 2])
>>> ds = ds.repeat().apply(
...     tf.contrib.data.batch_and_drop_remainder(params['batch-size']))
>>> ds

 <_RestructuredDataset shapes: (32, 3), types: tf.int32>

顾名思义，这种方法的一个缺点就是会在数据集的结尾丢弃任何的未满批次.对于用于训练的无限重复数据集，这是可以接收的，但是你如果想要按一个具体的循环数训练，则会出现问题。

要进行一轮准确的运算，你可以通过手动填充批次的长度，并在创建 tf.metrics 时将条目设置为零权重来解决这一问题。

数据集

输入管道性能指南

tf.data.Dataset.from_tensor_slices

tf.data.TFRecordDataset

tf.data.Dataset.cache

性能指南

同样重要的是，无论使用哪种类型的读取器，都要使用构造函数的 buffer_size 参数启用缓冲。此参数以字节为单位指定。建议使用几 MB（ buffer_size=8*1024*1024 ），以便在需要时提供数据。

TPU 示例仓库下包含一个用于下载 ImageNet 数据集并将其转换为适当格式的 脚本 。

这与仓库中包含的 ImageNet 模型 一起演示了所有的最佳实践。

下一步

有关如何实际设置和运行 Cloud TPU 的详细信息，请参看：

• Google Cloud TPU 文档

这篇文章也不能包含所有。关于如何使 Cloud TPU 兼容模型的更多细节的最佳来源是发布在以下文章中的实例模型：

• TPU 示例仓库

有关优化 TensorFlow 代码以提高性能的更多信息，请参看：

• 性能

如果您发现本页面存在错误或可以改进，请 点击此处 帮助我们改进。