Sam Witteveen 的客座文章

您知道Colab 包括为训练模型选择免费 Cloud TPU 的功能吗？没错，您可以独享整个 TPU，在笔记本中使用！从 TensorFlow 1.11 开始，您就可以使用 TPU 训练 Keras 模型。

在本文中，让我们看看您需要对代码进行哪些更改才能在 TPU 上训练 Keras 模型。请注意，随着今年晚些时候 TensorFlow 2.0 的发布，其中一些内容可能会进一步简化，但我认为分享这些技巧会有所帮助，以防您想现在就尝试一下。

tf.data

让您的模型在 Cloud TPU 上运行的最大部分可能是设置正确的 数据管道。在构建要在 TPU 上运行的模型时，挑战在于性能瓶颈通常不再是在模型加速方面，而是在向模型馈送数据的管道方面。

我们要确保在训练模型时不会让 TPU 数据不足。为此，我们使用tf.data。让我们快速回顾一下我们需要什么。

首先，我们需要一个输入函数，它可以接收我们的数据并对其进行切片、混洗和批处理。由于图是使用XLA 编译的，因此我们需要预先指定张量的形状。这使得确保 TPU 每次都能获得相同大小的批次变得很重要。为此，我们使用参数“drop_remainder = True”，以便丢弃任何末尾的偏批次，而不是将其作为不同的批次大小呈现给模型。

def train_input_fn(batch_size=1024):
    # Convert the inputs to a Dataset.
    dataset = tf.data.Dataset.from_tensor_slices((x_train,y_train))
# Shuffle, repeat, and batch the examples.
    dataset = dataset.cache()
    dataset = dataset.shuffle(1000, reshuffle_each_iteration=True)
    dataset = dataset.repeat()
    dataset = dataset.batch(batch_size, drop_remainder=True)
# Return the dataset.
    return dataset

模型的更改

一旦我们设置了 tf.data 管道，我们只需对模型进行一些关键更改，使其与 TPU 兼容。

其中第一个是获取 TPU 的地址，以便我们稍后将其传递给分布式策略。这可以通过在笔记本的顶部运行以下代码来完成，以检查是否有 TPU 附加到 VM，如果有，则打印 TPU 的地址。

获取 TPU 地址

try:
 device_name = os.environ[‘COLAB_TPU_ADDR’]
 TPU_ADDRESS = ‘grpc://’ + device_name
 print(‘Found TPU at: {}’.format(TPU_ADDRESS))
except KeyError:
 print(‘TPU not found’)

接下来，我们需要选择用于模型的优化器。TensorFlow 优化器目前比 Keras 优化器支持得更好。在这里，您可以看到我们选择使用 TensorFlow Adam 优化器。

# Use a tf optimizer rather than a Keras one for now
opt = tf.train.AdamOptimizer(learning_rate)
model.compile(
      optimizer=opt,
      loss='categorical_crossentropy',
      metrics=['acc'])

最后，我们需要将 Keras 模型转换为 TPU 模型。目前，我们通过使用 keras_to_tpu 函数并传入一个分布式策略来实现。（最终，此步骤将移入编译函数中，您只需将分布式策略传递到 Keras 模型的编译函数中。）

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))

完成此操作后，您应该会看到类似于此的输出，显示 TPU 已准备好开始训练。我们还可以看到 TPU 设备的详细信息。

keras_to_tpu 函数的输出

批次大小

在 GPU 上，我们总是希望使批次大小足够大，以便同时使用尽可能多的 Cuda 内核。同样，在 TPU 上，我们希望确保我们的批次大小足够大，以充分利用 8 个内核及其每个 systolic 阵列。

systolic 阵列在目前可在 Colab 中访问的 Cloud TPU v2 中为 128 x 128。使用此信息，我们可以确定每个内核的批次大小应该是 128 的倍数。实现此目的的最简单方法是使用 1024 的全局批次大小（每个 8 个内核 128）。

除了更改批次大小外，您可能还想调整学习率以适应更大的批次大小。

训练

我们现在可以将 tf.data 输入函数作为数据集传递，并像通常训练 Keras 模型一样训练您的 TPU 模型。

tpu_model.fit(
    train_input_fn,
    steps_per_epoch = 60,
    epochs=10,
)

展望未来 — 分布式策略

在 TensorFlow 和 tf.keras 的未来几个版本中，我们预计将引入分布式策略，这将使使用 TPU 训练模型变得更加容易。使用此新 API，我们只需在编译模型时传入分布式策略即可。

结论

为了回顾一下在 Cloud TPU 上运行模型需要进行的关键更改，它们是

• 查找您的 TPU 地址
• 为固定张量大小设置您的模型和 tf.data
• 将 Keras 模型转换为 TPU Keras 模型
• 选择正确的批次大小
• 使用 tf.data 为您的模型提供数据

现在还处于早期阶段，但这是一种尝试使用 TPU 的激动人心的方式！有关使用 TensorFlow 与 TPU 的更多信息，请查看Cloud TPU 性能指南 和官方TPU 示例。

欢迎在 Colab 中尝试我编写的示例此处。