如何利用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/
2008墨西哥摩托车展会
测量实际丛林环境下的电波传播路径损耗及天线效应
2020年中国信息产业收入将增至26.2万亿元 九大发展重点
开发用于重型履带式汽车的双测功机道路负载仿真器
国内首部3D电视标准已完成定稿
如何利用Google Colab的云TPU加速Keras模型训练
单片机学习笔记
装机量增长毛利率下滑 LFP企业“承压”
Android Oreo 内置回滚保护，禁止操作系统降级
数字电路设计之同步时序逻辑电路
第四届智慧城市照明发展论坛成功召开
化工厂5G+蓝牙+LoRa室内人员定位系统解决方案
发现量子反常霍尔效应是中国对人类科学重要贡献
简述新能源汽车充电设备设计中的常见问题
独石电容和电解电容的区别
区块链抵押贷款会带来怎样的影响
玺尔盾智能锁，一款没有指纹也能用的智能锁
给人工智能贴标的新“富士康工人”
共享单车之后还有雨伞、充电宝、床铺 难道共享睡眠舱也要火吗？
华为完成低频2.6GHz频段下5G基站的新空口测试
我以前都在单张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/
2008墨西哥摩托车展会
测量实际丛林环境下的电波传播路径损耗及天线效应
2020年中国信息产业收入将增至26.2万亿元 九大发展重点
开发用于重型履带式汽车的双测功机道路负载仿真器
国内首部3D电视标准已完成定稿
如何利用Google Colab的云TPU加速Keras模型训练
单片机学习笔记
装机量增长毛利率下滑 LFP企业“承压”
Android Oreo 内置回滚保护，禁止操作系统降级
数字电路设计之同步时序逻辑电路
第四届智慧城市照明发展论坛成功召开
化工厂5G+蓝牙+LoRa室内人员定位系统解决方案
发现量子反常霍尔效应是中国对人类科学重要贡献
简述新能源汽车充电设备设计中的常见问题
独石电容和电解电容的区别
区块链抵押贷款会带来怎样的影响
玺尔盾智能锁，一款没有指纹也能用的智能锁
给人工智能贴标的新“富士康工人”
共享单车之后还有雨伞、充电宝、床铺 难道共享睡眠舱也要火吗？
华为完成低频2.6GHz频段下5G基站的新空口测试