如何利用Google Colab的云TPU加速Keras模型训练

编者按：AI软件开发者Chengwei Zhang介绍了如何利用Google Colab的云TPU加速Keras模型训练。

我以前都在单张GTX 1070显卡（8.18 TFlops）上训练自己的模型，而且还挺满足的。后来Google的Colab开放了免费的Tesla K80显卡（12GB 显存，8.73 TFlops），最近又提供了免费的TPU（180 TFlops）。这篇教程将简要介绍如何将现有的Keras模型转换为TPU模型，然后在Colab上训练。和在我的GTX1070上训练相比，免费的TPU能够加速20倍。

我们将创建一个容易理解但训练起来足够复杂的Keras模型，让TPU热乎热乎。训练一个LSTM模型，进行IMDB情感分类任务，可能是一个很不错的例子，因为相比密集层和卷积层，训练LSTM对算力要求更高。

工作流概览：

创建Keras模型，输入采用固定的batch_size

转换Keras模型为TPU模型

以batch_size * 8训练TPU模型，并保存权重至文件

创建一个结构相同但输入batch大小可变的Keras模型，用于推理

加载模型权重

基于推理模型进行预测

在阅读本文的同时，你可以上手试验相应的Colab Jupyter notebook：https://colab.research.google.com/drive/1QZf1WeX3EQqBLeFeT4utFKBqq-ogG1FN

闲话少叙，让我们开始吧。

首先在Colab运行时激活TPU：

固定输入batch尺寸

大多数情况下，CPU和GPU上对输入形状没有限制，但XLA/TPU环境下强制使用固定的形状和batch尺寸。

云TPU包含8个TPU核，每个核都作为独立的处理单元运作。如果没有用上全部8个核心，那就没有充分利用TPU。为了充分加速训练，相比在单GPU上训练的同样的模型，我们可以选择较大的batch尺寸。总batch尺寸定为1024（每个核心128）一般是一个不错的起点。

万一你要训练一个较大的模型，batch尺寸太大了，那就慢慢降低batch尺寸，直到TPU的内存放得下为止。只需确保总batch尺寸是64的倍数（每核心的batch尺寸应该是8的倍数）。

值得一提的是，当batch尺寸较大时，一般来说增加优化算法的学习率以更快收敛的做法是安全的。详情参见Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour这篇论文（ arXiv:1706.02677）。

Keras允许通过参数batch_size设定输入层的batch尺寸。注意我们将模型定义为一个接受batch_size参数的函数，这样我们之后可以很方便地创建在CPU或GPU上运行的模型，这些模型接受可变batch尺寸的输入。

import tensorflow as tf

from tensorflow.python.keras.layers importInput, LSTM, Bidirectional, Dense, Embedding

def make_model(batch_size=None):

source = Input(shape=(maxlen,), batch_size=batch_size,

dtype=tf.int32, name='Input')

embedding = Embedding(input_dim=max_features,

output_dim=128, name='Embedding')(source)

lstm = LSTM(32, name='LSTM')(embedding)

predicted_var = Dense(1, activation='sigmoid', name='Output')(lstm)

model = tf.keras.Model(inputs=[source], outputs=[predicted_var])

model.compile(

optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),

loss='binary_crossentropy',

metrics=['acc'])

return model

training_model = make_model(batch_size=128)

另外，我们这里用了tf.train.Optimizer而不是标准的Keras优化器，因为TPU对Keras优化器的支持还处于实验阶段。

转换Keras模型至TPU模型

tf.contrib.tpu.keras_to_tpu_model函数可以转换tf.keras模型至等价的TPU版本。

import os

import tensorflow as tf

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']

tf.logging.set_verbosity(tf.logging.INFO)

tpu_model = tf.contrib.tpu.keras_to_tpu_model(

training_model,

strategy=tf.contrib.tpu.TPUDistributionStrategy(

tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

然后我们训练模型，保存权重，并评估模型。注意batch_size设定为模型输入batch_size的8倍，因为输入样本在8个TPU核心上均匀分布。

history = tpu_model.fit(x_train, y_train,

epochs=20,

batch_size=128 * 8,

validation_split=0.2)

tpu_model.save_weights('./tpu_model.h5', overwrite=True)

tpu_model.evaluate(x_test, y_test, batch_size=128 * 8)

我比较了单GTX1070显卡（在我的Windows电脑上本地运行）上和Colab的TPU上的训练速度，结果如下。

GPU和TPU的输入batch尺寸均为128.

GPU：179秒每epoch。20个epoch后达到了76.9%的验证精确度，共计3600秒。

TPU：5秒每epoch，第一个epoch除外（49秒）。20个epoch后达到了95.2%的验证精确度，共计150秒。

20个epoch后TPU的验证精确度高于在GPU上的表现，可能是因为TPU上同时训练8个batch的样本（每个batch大小为128）。

译者注：在Tesla K80上训练20个epoch后的验证精确度为86.3%（耗时6004秒）。使用TPU单核心训练（tf.contrib.tpu.TPUDistributionStrategy函数加上using_single_core=True参数）20个epoch后达到了99.8%的验证精确度。将模型的batch尺寸改为16（128/8）后，TPU上训练20个epoch后达到了99.8%的验证精确度（因为batch尺寸改变，训练时间延长了，约377秒）。这样看起来TensorFlow在TPU上的实现可能有些问题。

在CPU上推理

得到模型权重后，我们可以像平时一样加载权重，然后在其他设备（比如CPU或GPU）上做出预测。我们同时想要推理模型接受可变的输入batch尺寸，如前所述，只需给make_model()函数指定一个参数即可。

inferencing_model = make_model(batch_size=None)

inferencing_model.load_weights('./tpu_model.h5')

inferencing_model.summary()

summary()方法的输出表明推理模型现在可以接受可变输入样本数目：

_________________________________________________________________

Layer (type) OutputShapeParam#

=================================================================

Input (InputLayer) (None, 500) 0

_________________________________________________________________

Embedding (Embedding) (None, 500, 128) 1280000

_________________________________________________________________

LSTM (LSTM) (None, 32) 20608

_________________________________________________________________

Output (Dense) (None, 1) 33

=================================================================

接下来我们就可以在推理模型上调用标准的fit()、evaluate()函数。

inferencing_model.evaluate(x_test, y_test)

我们的模型在测试集上的精确度为82.4%

25000/25000 [==============================] - 83s3ms/step

[0.6060782189846039, 0.824]

译者注：相比测试损失和测试精确度，验证损失和验证精确度太高了。其他两个在TPU上训练的模型（单核训练和不同batch大小）也出现了类似的现象，进一步加大了TensorFlow在TPU上的实现有问题的嫌疑。

最后，你可以下载模型权重到本地，以便以后在其他地方使用。

from google.colab import files

files.download('./tpu_model.h5')

结语和进一步阅读

这篇教程简要介绍了如何利用Google Colab的免费云TPU资源，加速Keras模型的训练。

云TPU文档：https://cloud.google.com/tpu/docs/

云TPU性能指南：https://cloud.google.com/tpu/docs/performance-guide

云TPU排错指南：https://cloud.google.com/tpu/docs/troubleshooting

XLA概览：https://www.tensorflow.org/performance/xla/