在 Colab 中使用 TPU 训练 tf.keras

文 / 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 是否已附加至虚拟机，如果是，则打印 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’)

接下来，我们需要选择用于此模型的优化器。与 Keras 优化器相比，TensorFlow 优化器目前的支持程度更好。在下面您可以看到，我们已选择使用 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 个核心及其每个脉动阵列。

这些脉动阵列在 Cloud TPU v2 中的大小是 128 x 128，目前可以在 Colab 中获取。借助此信息，我们可以确定，对每个核心而言，批次大小应为 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 示例（https://cloud.google.com/tpu/docs/performance-guide）。

欢迎试用我在 Colab 中编写的 示例（https://colab.research.google.com/github/samwit/TPU/blob/master/Guide_to_tf_keras_TPU_MNIST.ipynb）。

更多 AI 相关阅读：