如何利用Google Colab的云TPU加速Keras模型訓(xùn)練

編者按：AI軟件開發(fā)者Chengwei Zhang介紹了如何利用Google Colab的云TPU加速Keras模型訓(xùn)練。

我以前都在單張GTX 1070顯卡（8.18 TFlops）上訓(xùn)練自己的模型，而且還挺滿足的。后來Google的Colab開放了免費(fèi)的Tesla K80顯卡（12GB 顯存，8.73 TFlops），最近又提供了免費(fèi)的TPU（180 TFlops）。這篇教程將簡要介紹如何將現(xiàn)有的Keras模型轉(zhuǎn)換為TPU模型，然后在Colab上訓(xùn)練。和在我的GTX1070上訓(xùn)練相比，免費(fèi)的TPU能夠加速20倍。

我們將創(chuàng)建一個(gè)容易理解但訓(xùn)練起來足夠復(fù)雜的Keras模型，讓TPU熱乎熱乎。訓(xùn)練一個(gè)LSTM模型，進(jìn)行IMDB情感分類任務(wù)，可能是一個(gè)很不錯(cuò)的例子，因?yàn)橄啾让芗瘜雍途矸e層，訓(xùn)練LSTM對(duì)算力要求更高。

工作流概覽：

創(chuàng)建Keras模型，輸入采用固定的batch_size

轉(zhuǎn)換Keras模型為TPU模型

以batch_size * 8訓(xùn)練TPU模型，并保存權(quán)重至文件

創(chuàng)建一個(gè)結(jié)構(gòu)相同但輸入batch大小可變的Keras模型，用于推理

加載模型權(quán)重

基于推理模型進(jìn)行預(yù)測(cè)

在閱讀本文的同時(shí)，你可以上手試驗(yàn)相應(yīng)的Colab Jupyter notebook：https://colab.research.google.com/drive/1QZf1WeX3EQqBLeFeT4utFKBqq-ogG1FN

閑話少敘，讓我們開始吧。

首先在Colab運(yùn)行時(shí)激活TPU：

固定輸入batch尺寸

大多數(shù)情況下，CPU和GPU上對(duì)輸入形狀沒有限制，但XLA/TPU環(huán)境下強(qiáng)制使用固定的形狀和batch尺寸。

云TPU包含8個(gè)TPU核，每個(gè)核都作為獨(dú)立的處理單元運(yùn)作。如果沒有用上全部8個(gè)核心，那就沒有充分利用TPU。為了充分加速訓(xùn)練，相比在單GPU上訓(xùn)練的同樣的模型，我們可以選擇較大的batch尺寸。總batch尺寸定為1024（每個(gè)核心128）一般是一個(gè)不錯(cuò)的起點(diǎn)。

萬一你要訓(xùn)練一個(gè)較大的模型，batch尺寸太大了，那就慢慢降低batch尺寸，直到TPU的內(nèi)存放得下為止。只需確?？俠atch尺寸是64的倍數(shù)（每核心的batch尺寸應(yīng)該是8的倍數(shù)）。

值得一提的是，當(dāng)batch尺寸較大時(shí)，一般來說增加優(yōu)化算法的學(xué)習(xí)率以更快收斂的做法是安全的。詳情參見Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour這篇論文（ arXiv:1706.02677）。

Keras允許通過參數(shù)batch_size設(shè)定輸入層的batch尺寸。注意我們將模型定義為一個(gè)接受batch_size參數(shù)的函數(shù)，這樣我們之后可以很方便地創(chuàng)建在CPU或GPU上運(yùn)行的模型，這些模型接受可變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而不是標(biāo)準(zhǔn)的Keras優(yōu)化器，因?yàn)門PU對(duì)Keras優(yōu)化器的支持還處于實(shí)驗(yàn)階段。

轉(zhuǎn)換Keras模型至TPU模型

tf.contrib.tpu.keras_to_tpu_model函數(shù)可以轉(zhuǎn)換tf.keras模型至等價(jià)的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)))

然后我們訓(xùn)練模型，保存權(quán)重，并評(píng)估模型。注意batch_size設(shè)定為模型輸入batch_size的8倍，因?yàn)檩斎霕颖驹?個(gè)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電腦上本地運(yùn)行）上和Colab的TPU上的訓(xùn)練速度，結(jié)果如下。

GPU和TPU的輸入batch尺寸均為128.

GPU：179秒每epoch。20個(gè)epoch后達(dá)到了76.9%的驗(yàn)證精確度，共計(jì)3600秒。

TPU：5秒每epoch，第一個(gè)epoch除外（49秒）。20個(gè)epoch后達(dá)到了95.2%的驗(yàn)證精確度，共計(jì)150秒。

20個(gè)epoch后TPU的驗(yàn)證精確度高于在GPU上的表現(xiàn)，可能是因?yàn)門PU上同時(shí)訓(xùn)練8個(gè)batch的樣本（每個(gè)batch大小為128）。

譯者注：在Tesla K80上訓(xùn)練20個(gè)epoch后的驗(yàn)證精確度為86.3%（耗時(shí)6004秒）。使用TPU單核心訓(xùn)練（tf.contrib.tpu.TPUDistributionStrategy函數(shù)加上using_single_core=True參數(shù)）20個(gè)epoch后達(dá)到了99.8%的驗(yàn)證精確度。將模型的batch尺寸改為16（128/8）后，TPU上訓(xùn)練20個(gè)epoch后達(dá)到了99.8%的驗(yàn)證精確度（因?yàn)閎atch尺寸改變，訓(xùn)練時(shí)間延長了，約377秒）。這樣看起來TensorFlow在TPU上的實(shí)現(xiàn)可能有些問題。

在CPU上推理

得到模型權(quán)重后，我們可以像平時(shí)一樣加載權(quán)重，然后在其他設(shè)備（比如CPU或GPU）上做出預(yù)測(cè)。我們同時(shí)想要推理模型接受可變的輸入batch尺寸，如前所述，只需給make_model()函數(shù)指定一個(gè)參數(shù)即可。

inferencing_model = make_model(batch_size=None)

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

inferencing_model.summary()

summary()方法的輸出表明推理模型現(xiàn)在可以接受可變輸入樣本數(shù)目：

_________________________________________________________________

Layer (type) OutputShapeParam#

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

Input (InputLayer) (None, 500) 0

_________________________________________________________________

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

_________________________________________________________________

LSTM (LSTM) (None, 32) 20608

_________________________________________________________________

Output (Dense) (None, 1) 33

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

接下來我們就可以在推理模型上調(diào)用標(biāo)準(zhǔn)的fit()、evaluate()函數(shù)。

inferencing_model.evaluate(x_test, y_test)

我們的模型在測(cè)試集上的精確度為82.4%

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

[0.6060782189846039, 0.824]

譯者注：相比測(cè)試損失和測(cè)試精確度，驗(yàn)證損失和驗(yàn)證精確度太高了。其他兩個(gè)在TPU上訓(xùn)練的模型（單核訓(xùn)練和不同batch大?。┮渤霈F(xiàn)了類似的現(xiàn)象，進(jìn)一步加大了TensorFlow在TPU上的實(shí)現(xiàn)有問題的嫌疑。

最后，你可以下載模型權(quán)重到本地，以便以后在其他地方使用。

from google.colab import files

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

結(jié)語和進(jìn)一步閱讀

這篇教程簡要介紹了如何利用Google Colab的免費(fèi)云TPU資源，加速Keras模型的訓(xùn)練。

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

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

云TPU排錯(cuò)指南：https://cloud.google.com/tpu/docs/troubleshooting

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