本文介紹了如何利用 Google Colab 上的免費(fèi) Cloud TPU 資源更快地訓(xùn)練 Keras 模型,。

很長一段時(shí)間以來，我在單個(gè) GTX 1070 顯卡上訓(xùn)練模型,，其單精度大約為 8.18 TFlops,。后來谷歌在 Colab 上啟用了免費(fèi)的 Tesla K80 GPU，配備 12GB 內(nèi)存,，且速度稍有增加,，為 8.73 TFlops。最近,，Colab 的運(yùn)行時(shí)類型選擇器中出現(xiàn)了 Cloud TPU 選項(xiàng),，其浮點(diǎn)計(jì)算能力為 180 TFlops。

本文將介紹如何在 Colab 上使用 TPU 訓(xùn)練已有的 Keras 模型,，其訓(xùn)練速度是在 GTX 1070 上訓(xùn)練速度的 20 倍,。

我們首先構(gòu)建一個(gè)易于理解但訓(xùn)練過程比較復(fù)雜的 Keras 模型，以便「預(yù)熱」Cloud TPU,。在 IMDB 情感分類任務(wù)上訓(xùn)練 LSTM 模型是個(gè)不錯(cuò)的選擇,，因?yàn)?LSTM 的計(jì)算成本比密集和卷積等層高。

流程如下所示：

• 構(gòu)建一個(gè) Keras 模型,，可使靜態(tài)輸入 batch_size 在函數(shù)式 API 中進(jìn)行訓(xùn)練,。

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

• 使用靜態(tài) batch_size * 8 訓(xùn)練 TPU 模型,，并將權(quán)重保存到文件,。

• 構(gòu)建結(jié)構(gòu)相同但輸入批大小可變的 Keras 模型，用于執(zhí)行推理,。

• 加載模型權(quán)重,。

• 使用推理模型進(jìn)行預(yù)測。

讀者閱讀本文時(shí),，可以使用 Colab Jupyter notebook Keras_LSTM_TPU.ipynb（https://colab.research.google.com/drive/1QZf1WeX3EQqBLeFeT4utFKBqq-ogG1FN）進(jìn)行試驗(yàn),。

首先，按照下圖的說明在 Colab 運(yùn)行時(shí)選項(xiàng)中選擇激活 TPU,。

激活 TPU

靜態(tài)輸入 Batch Size

在 CPU 和 GPU 上運(yùn)行的輸入管道大多沒有靜態(tài)形狀的要求,，而在 XLA/TPU 環(huán)境中，則對靜態(tài)形狀和 batch size 有要求,。

Could TPU 包含 8 個(gè)可作為獨(dú)立處理單元運(yùn)行的 TPU 核心,。只有八個(gè)核心全部工作，TPU 才算被充分利用,。為通過向量化充分提高訓(xùn)練速度,，我們可以選擇比在單個(gè) GPU 上訓(xùn)練相同模型時(shí)更大的 batch size,。最開始最好設(shè)定總 batch size 為 1024（每個(gè)核心 128 個(gè)）。

如果你要訓(xùn)練的 batch size 過大,，可以慢慢減小 batch size,，直到它適合 TPU 內(nèi)存，只需確?？偟?batch size 為 64 的倍數(shù)即可（每個(gè)核心的 batch size 大小應(yīng)為 8 的倍數(shù)）,。

使用較大的 batch size 進(jìn)行訓(xùn)練也同樣有價(jià)值：通常可以穩(wěn)定地提高優(yōu)化器的學(xué)習(xí)率,，以實(shí)現(xiàn)更快的收斂,。（參考論文：https:///pdf/1706.02677.pdf）

在 Keras 中，要定義靜態(tài) batch size,，我們需使用其函數(shù)式 API,，然后為 Input 層指定 batch_size 參數(shù)。請注意,，模型在一個(gè)帶有 batch_size 參數(shù)的函數(shù)中構(gòu)建,，這樣方便我們再回來為 CPU 或 GPU 上的推理運(yùn)行創(chuàng)建另一個(gè)模型,，該模型采用可變的輸入 batch size,。

此外，使用 tf.train.Optimizer,，而不是標(biāo)準(zhǔn)的 Keras 優(yōu)化器,，因?yàn)?Keras 優(yōu)化器對 TPU 而言還處于試驗(yàn)階段。

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

tf.contrib.tpu.keras_to_tpu_model 函數(shù)將 tf.keras 模型轉(zhuǎn)換為同等的 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)))

然后使用標(biāo)準(zhǔn)的 Keras 方法來訓(xùn)練,、保存權(quán)重并評估模型。請注意,，batch_size 設(shè)置為模型輸入 batch_size 的八倍,，這是為了使輸入樣本在 8 個(gè) TPU 核心上均勻分布并運(yùn)行。

我設(shè)置了一個(gè)實(shí)驗(yàn),，比較在 Windows PC 上使用單個(gè) GTX1070 和在 Colab 上運(yùn)行 TPU 的訓(xùn)練速度,，結(jié)果如下。

GPU 和 TPU 都將輸入 batch size 設(shè)為 128,，

• GPU：每個(gè) epoch 需要 179 秒,。20 個(gè) epoch 后驗(yàn)證準(zhǔn)確率達(dá)到 76.9％，總計(jì) 3600 秒,。

• TPU：每個(gè) epoch 需要 5 秒,，第一個(gè) epoch 除外（需 49 秒）。20 個(gè) epoch 后驗(yàn)證準(zhǔn)確率達(dá)到 95.2％,，總計(jì) 150 秒,。

20 個(gè) epoch 后,，TPU 上訓(xùn)練模型的驗(yàn)證準(zhǔn)確率高于 GPU，這可能是由于在 GPU 上一次訓(xùn)練 8 個(gè) batch,，每個(gè) batch 都有 128 個(gè)樣本,。

在 CPU 上執(zhí)行推理

一旦我們獲得模型權(quán)重，就可以像往常一樣加載它,，并在 CPU 或 GPU 等其他設(shè)備上執(zhí)行預(yù)測,。我們還希望推理模型接受靈活的輸入 batch size，這可以使用之前的 make_model() 函數(shù)來實(shí)現(xiàn),。

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

可以看到推理模型現(xiàn)在采用了可變的輸入樣本,。

然后，你可以使用標(biāo)準(zhǔn) fit(),、evaluate() 函數(shù)與推理模型,。

結(jié)論

本快速教程介紹了如何利用 Google Colab 上的免費(fèi) Cloud TPU 資源更快地訓(xùn)練 Keras 模型。

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