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

电子说

1.3w人已加入

描述

编者按: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/

打开APP阅读更多精彩内容
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉

全部0条评论

快来发表一下你的评论吧 !

×
20
完善资料,
赚取积分