如何利用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尺寸?？俠atch尺寸定為1024（每個核心128）一般是一個不錯的起點。

萬一你要訓練一個較大的模型，batch尺寸太大了，那就慢慢降低batch尺寸，直到TPU的內存放得下為止。只需確?？俠atch尺寸是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大?。┮渤霈F了類似的現象，進一步加大了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/