电子说
编者按: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/
全部0条评论
快来发表一下你的评论吧 !