用免費TPU訓練Keras模型，速度還能提高20倍！

選自KDnuggets

作者：Chengwei Zhang

機器之心編譯

參與：高璇、路

本文介紹了如何利用 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。

此外，使用 tf.train.Optimizer，而不是標準的 Keras 優化器，因為 Keras 優化器對 TPU 而言還處於試驗階段。

將 Keras 模型轉換為 TPU 模型

然後使用標準的 Keras 方法來訓練、保存權重並評估模型。請注意，batch_size 設置為模型輸入 batch_size 的八倍，這是為了使輸入樣本在 8 個 TPU 核心上均勻分布並運行。

我設置了一個實驗，比較在 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() 函數來實現。

可以看到推理模型現在採用了可變的輸入樣本。

然後，你可以使用標準 fit()、evaluate() 函數與推理模型。

結論

本快速教程介紹了如何利用 Google Colab 上的免費 Cloud TPU 資源更快地訓練 Keras 模型。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！