Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

参与高璇 路雪

用免费TPU训练Keras模型,速度还能提高20倍!

本文介绍了如何利用 Google Colab 上的免费 Cloud TPU 资源更快地训练 Keras 模型。

很长一段时间以来,我在单个 GTX 1070 显卡上训练模型,其单精度大约为 8.18 TFlops。后来谷歌在 Colab 上启用了免费的 Tesla K80 GPU,配备 12GB 内存,且速度稍有增加,为 8.73 TFlops。最近,Colab 的运行时类型选择器中出现了 Cloud TPU 选项,其浮点计算能力为 180 TFlops。

本文将介绍如何在 Colab 上使用 TPU 训练已有的 Keras 模型,其训练速度是在 GTX 1070 上训练速度的 20 倍。

我们首先构建一个易于理解但训练过程比较复杂的 Keras 模型,以便「预热」Cloud TPU。在 IMDB 情感分类任务上训练 LSTM 模型是个不错的选择,因为 LSTM 的计算成本比密集和卷积等层高。

流程如下所示:

  • 构建一个 Keras 模型,可使静态输入 batch_size 在函数式 API 中进行训练。

  • 将 Keras 模型转换为 TPU 模型。

  • 使用静态 batch_size * 8 训练 TPU 模型,并将权重保存到文件。

  • 构建结构相同但输入批大小可变的 Keras 模型,用于执行推理。

  • 加载模型权重

  • 使用推理模型进行预测。

读者阅读本文时,可以使用 Colab Jupyter notebook Keras_LSTM_TPU.ipynb(https://colab.research.google.com/drive/1QZf1WeX3EQqBLeFeT4utFKBqq-ogG1FN)进行试验。

首先,按照下图的说明在 Colab 运行时选项中选择激活 TPU。

激活 TPU

静态输入 Batch Size

在 CPU 和 GPU 上运行的输入管道大多没有静态形状的要求,而在 XLA/TPU 环境中,则对静态形状和 batch size 有要求。

Could TPU 包含 8 个可作为独立处理单元运行的 TPU 核心。只有八个核心全部工作,TPU 才算被充分利用。为通过向量化充分提高训练速度,我们可以选择比在单个 GPU 上训练相同模型时更大的 batch size。最开始最好设定总 batch size 为 1024(每个核心 128 个)。

如果你要训练的 batch size 过大,可以慢慢减小 batch size,直到它适合 TPU 内存,只需确保总的 batch size 为 64 的倍数即可(每个核心的 batch size 大小应为 8 的倍数)。

使用较大的 batch size 进行训练也同样有价值:通常可以稳定地提高优化器学习率,以实现更快的收敛。(参考论文:https://arxiv.org/pdf/1706.02677.pdf)

在 Keras 中,要定义静态 batch size,我们需使用其函数式 API,然后为 Input 层指定 batch_size 参数。请注意,模型在一个带有 batch_size 参数的函数中构建,这样方便我们再回来为 CPU 或 GPU 上的推理运行创建另一个模型,该模型采用可变的输入 batch size。

import tensorflow as tf
from tensorflow.python.keras.layers import Input, 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 优化器,因为 Keras 优化器对 TPU 而言还处于试验阶段。

将 Keras 模型转换为 TPU 模型

tf.contrib.tpu.keras_to_tpu_model 函数将 tf.keras 模型转换为同等的 TPU 模型。

import os
import tensorflow as tf
# This address identifies the TPU we'll use when configuring TensorFlow.
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)))

然后使用标准的 Keras 方法来训练、保存权重并评估模型。请注意,batch_size 设置为模型输入 batch_size 的八倍,这是为了使输入样本在 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)

我设置了一个实验,比较在 Windows PC 上使用单个 GTX1070 和在 Colab 上运行 TPU 的训练速度,结果如下。

GPU 和 TPU 都将输入 batch size 设为 128,

  • GPU:每个 epoch 需要 179 秒。20 个 epoch 后验证准确率达到 76.9%,总计 3600 秒。

  • TPU:每个 epoch 需要 5 秒,第一个 epoch 除外(需 49 秒)。20 个 epoch 后验证准确率达到 95.2%,总计 150 秒。

20 个 epoch 后,TPU 上训练模型的验证准确率高于 GPU,这可能是由于在 GPU 上一次训练 8 个 batch,每个 batch 都有 128 个样本。

在 CPU 上执行推理

一旦我们获得模型权重,就可以像往常一样加载它,并在 CPU 或 GPU 等其他设备上执行预测。我们还希望推理模型接受灵活的输入 batch size,这可以使用之前的 make_model() 函数来实现。

inferencing_model = make_model(batch_size=None)
inferencing_model.load_weights('./tpu_model.h5')
inferencing_model.summary()

可以看到推理模型现在采用了可变的输入样本。

_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
Input (InputLayer) (None, 500) 0
_________________________________________________________________
Embedding (Embedding) (None, 500, 128) 1280000
_________________________________________________________________
LSTM (LSTM) (None, 32) 20608
_________________________________________________________________
Output (Dense) (None, 1) 33
=================================================================

然后,你可以使用标准 fit()、evaluate() 函数与推理模型。

结论

本快速教程介绍了如何利用 Google Colab 上的免费 Cloud TPU 资源更快地训练 Keras 模型。

原文链接:https://www.kdnuggets.com/2019/03/train-keras-model-20x-faster-tpu-free.html

工程谷歌TPUKeras
5
相关数据
权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

收敛技术

在数学,计算机科学和逻辑学中,收敛指的是不同的变换序列在有限的时间内达到一个结论(变换终止),并且得出的结论是独立于达到它的路径(他们是融合的)。 通俗来说,收敛通常是指在训练期间达到的一种状态,即经过一定次数的迭代之后,训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说,如果采用当前数据进行额外的训练将无法改进模型,模型即达到收敛状态。在深度学习中,损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变,暂时形成收敛的假象。

学习率技术

在使用不同优化器(例如随机梯度下降,Adam)神经网络相关训练中,学习速率作为一个超参数控制了权重更新的幅度,以及训练的速度和精度。学习速率太大容易导致目标(代价)函数波动较大从而难以找到最优,而弱学习速率设置太小,则会导致收敛过慢耗时太长

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

奇虎360机构

360公司成立于2005年8月,创始人周鸿祎 2011年3月30日在纽交所成功上市 2018年2月28日,回归A股上市,上证交易所(601360) 是中国第一大互联网安全公司,用户6.5亿,市场渗透率94.7% 中国第一大移动互联网安全公司,用户数超过8.5亿 中国领先的AIoT公司,将人工智能技术应用于智能生活、家庭安防、出行安全、儿童安全等多个领域

http://smart.360.cn/cleanrobot/
推荐文章
暂无评论
暂无评论~