Kaggle競賽硬體如何選擇？不差錢、追求速度，那就上TPU吧

選自towardsdatascience

作者：Paul Mooney

機器之心編譯

在每種機器學習算法背後，都是以數千兆赫頻率運算的硬體提供支持。你可能已經注意到，在設置 Kaggle Notebooks 計算環境時，有幾種處理器可供選擇，但究竟哪種才是最適合你的呢？本文分別比較了在訓練由 tf.keras 編寫的機器學習模型過程中，使用 CPU（英特爾至強處理器）、GPU（NVIDIA Tesla P100 顯卡）和 TPU（谷歌 TPU v3）的優劣勢。希望本文可以為 Kaggle 競賽準備者理解和選擇處理器提供幫助。

圖 1：在 Kaggle Notebook 中可以免費使用 CPU、GPU 和 TPU。

如何準備測試？

為了比較 CPU、GPU 和 TPU 完成常見數據科學任務時的性能，我們使用 tf_flowers 數據集訓練了一個卷積神經網絡，然後用三種不同的後端（GPU 是NVIDIA Tesla P100+2GHz、雙核的英特至強 CPU，13 GB RAM；TPU 是谷歌 8 核的 TPUv3+2GHz、四核的英特爾至強 CPU，16 GB RAM）將相同的代碼運行了三遍。所附的 Kaggle 教程筆記（tutorial notebook）顯示，在一些情況下，TPU 性能最好。

例如：

用分片文檔數據集（例如 TFRecord）；

用 tf.data API 將訓練數據傳到 TPU；

batch size 較大時（如 batch_size=128）。

在工作流程中加入這些步驟，可以避免常見的 I/O 瓶頸，否則這樣的 I/O 瓶頸會妨礙 TPU 發揮全部功能。用戶可以訪問 Kaggle TPU 官方文檔，了解其他一些在 TPU 上運行代碼的優化技巧。

官網鏈接：https://www.kaggle.com/docs/tpu

硬體性能如何

這三種硬體之間最明顯的區別在於使用 tf.keras 訓練模型的速度。tf.keras 庫是最流行的機器學習框架之一，因為 tf.keras 可以更快、更方便地實驗新想法。如果寫代碼的時間短，則就可以撥出更多時間進行計算。如果等代碼跑完的時間短，也就可以有更多時間評估新想法了。在機器學習比賽中，tf.keras 和 TPU 是一對很強大的組合！

圖 2：tf.keras 概貌。(A)tf.keras 是最流行的實現深度學習模型的框架之一；(B) 在設計深度學習模型時，tf.keras 可以快速進行實驗和迭代。

在下圖 3 第一個實驗中，我們在三種硬體上運行了相同的代碼（官方教程筆記中的改進版），並且為了避免 CPU 和 GPU 出現記憶體不足的問題，batch size 也設置得很小（設為 16）。在這種情況下可以看出，訓練 Xception 模型時，TPU 比 CPU 快了約 100 倍，比 GPU 快了約 3.5 倍，這是因為 TPU 處理批大小很大的數據時效率更高。我們也試著將 batch size 增加到 128，結果發現 TPU 速度快了 2 倍，GPU 和 CPU 均出現了記憶體不足的問題。所以，與之前的實驗相比，TPU 訓練 Xception 模型的速度比 GPU 快 7 倍。

官方教程筆記：https://www.kaggle.com/mgornergoogle/five-flowers-with-keras-and-xception-on-tpu

圖 3：用 CPU、GPU 和 TPU 訓練 Xception 模型（12 個 epoch）所需時間對比。

如下圖 4 所示，模型訓練的加速情況也與模型類別有關，Xception 和 Vgg16 就比 ResNet50 表現更好。在這種邊界很大的情況下，模型訓練速度是 TPU 唯一超過 GPU 的地方。例如，當預測少量樣本時，TPU 比 CPU 快了約 3 倍，但又比 GPU 慢了約 3 倍（在某些情況，比如在 batch size 很大的情況下做預測，TPU 表現出色，但本實驗中不存在這樣的情況）。

圖 4：模型類型和任務類型不同，加速情況也不同。(A) 與 ResNet50 相比，Xception 和 Vgg16 的提速更為明顯。(B) 在諸如少量樣本預測等特定任務中，GPU 的表現比 TPU 好。

為了使結果更為充分，我們注意到 Yu Emma Wang 等人 [1] 開發了一個叫做 ParaDnn 的嚴格基準，這一基準可以用來比較不同硬體訓練機器學習模型時的性能。利用 Yu Emma Wang 等人 [1] 的方法，可以總結出：用 TPU 代替 GPU 時，參數化模型性能可以提升 1 到 10 倍，使用模型的性能也可以提升 3 到 6.8 倍（下圖 5）。分片數據、很大的數據批和大模型結合在一起時，TPU 是最佳選擇。

圖 5：Wang[1] 論文中的重要發現。當 batch size 和 CNN 模型都較大時，TPU 的性能最好。你可以在 Kumar[2] 和 Jouppi[3] 的論文中，找到其他基準實驗。

[1] https://arxiv.org/abs/1907.10701

訓練模型時的價格考量

雖然平等地比較了硬體，但它們在價格上有相當大的差異。TPU（谷歌 TPU v3 每小時 8.00 美元，GCP 上可以按需選擇的 TPU v2 每小時 4.50 美元）比 GPU（NVIDIA Tesla P100 每小時 1.46 美元）貴了大概五倍。雖然如此，如果你想優化成本，那還是應該選擇 TPU，因為在訓練相同模型的情況下，TPU 的訓練速度至少快了 5 倍。

當數據以分片格式儲存在 GCS bucket，然後以大 batch size 傳遞到 TPU 時，模型訓練會提速約 5 倍，因此建議熟悉 tf.data API 的用戶使用 TPU。

有些機器學習實踐者優先考慮模型訓練時間，而不是模型訓練成本。因此，對於想盡快完成模型訓練的人來說，TPU 是最佳選擇。在訓練模型上花更少的時間，就可以花更多時間考慮新想法。但請不要一味聽信本文內容——你可以通過 Kaggle Notebooks 在 CPU、GPU 和 TPU 上免費運行代碼並評估結果。Kaggle 用戶們已經成功在 TPU 上運行了文本數據，並從中獲得了許多樂趣。

下面這篇文章的作者描述了是如何用 TPU 訓練 BERT 模型，並在最近的 Kaggle 競賽中贏得了 8000 美元獎金（二等獎）的。

文章鏈接：https://www.kaggle.com/c/tensorflow2-question-answering/discussion/127333

該選哪種硬體呢？

總之，如果有通用性和大記憶體需求的話，我們建議你使用 CPU。當你想加快各種數據科學流程時，GPU 是很好的替代品。如果想要盡可能快地完成模型訓練的話，TPU 則是最佳選擇。

針對所用硬體來優化代碼，可以得到更好的結果。我們認為，對分別針對 GPU 代碼和 TPU 代碼的運行時（runtime）進行比較也很有意思。例如，用像 RAPIDS.ai 這樣的 GPU 加速庫訓練梯度提升模型，再用像 tf.keras 這樣的 TPU 加速庫訓練深度學習模型，比較二者的訓練時間，這也是很有意思的。

那麽，訓練準確的機器學習模型最少需要多長時間？一天內可以評價多少不同的想法？結合 tf.keras，TPU 讓機器學習從業人員寫代碼的時間變短了，等代碼跑完的時間也變短了，這樣就可以留出更多時間評估新想法，並在 Kaggle 大賽中提升自己作品的性能。

原文鏈接：https://towardsdatascience.com/when-to-use-cpus-vs-gpus-vs-tpus-in-a-kaggle-competition-9af708a8c3eb

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

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

加入機器之心（全職記者 / 實習生）：[email protected]

投稿或尋求報導：content@jiqizhixin.com

廣告 & 商務合作：[email protected]