使用 TPU 加速訓練

當模型龐大, 例如: Efficientnet V2 L, 且資料數據集樣本多時, 我們用 Nvidia RTX 3090 具備了24GB VRAM 來訓練模型也產生 OOM 異常時, 當然可以進一步考慮多GPU卡, 或者 RTX A6000 有 48GB VRAM 也是一個選擇. 不過還有

使用 TPU 一個比較大的限制: 需要將數據集放置於 Google Cloud Storage, 通常還需要把數據集以 TF Records 格式儲存於 Google Cloud Storage 某 Bucket, 有些細節主要是錢 (要設置一個專案)與權(授權 TPU 讀TF Records ), 我們另外談.

另外一個問題是訓練程式在GPU跑得好好的, 直接拿來用是不行的, 通常只要套上樣板, 類似:

通常後面跟採用 GPU 訓練並無差異, 但是當我們訂製了自己的訓練時, 例如我們要進行 knowledge distiller 或者 noisy student 訓練時, 會遭遇問題, 損失函數原來為:

執行時將遭遇到 錯誤提示:

ValueError: Please use `tf.keras.losses.Reduction.SUM` or `tf.keras.losses.Reduction.NONE` for loss reduction when losses are used with `tf.distribute.Strategy` outside of the built-in training loops. You can implement `tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE` using global batch size like:

with strategy.scope():

loss_obj = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)

我們改為:

Under Construction