很長一段時間以來，我在單個 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