選自towardsdatascience
作者：Paul Mooney
機(jī)器之心編譯
在每種機(jī)器學(xué)習(xí)算法背后，都是以數(shù)千兆赫頻率運(yùn)算的硬件提供支持。你可能已經(jīng)注意到，在設(shè)置 Kaggle Notebooks 計(jì)算環(huán)境時(shí)，有幾種處理器可供選擇，但究竟哪種才是最適合你的呢？本文分別比較了在訓(xùn)練由 tf.keras 編寫(xiě)的機(jī)器學(xué)習(xí)模型過(guò)程中，使用 CPU（英特爾至強(qiáng)處理器）、GPU（英偉達(dá) Tesla P100 顯卡）和 TPU（谷歌 TPU v3）的優(yōu)劣勢(shì)。希望本文可以為 Kaggle 競(jìng)賽準(zhǔn)備者理解和選擇處理器提供幫助。
圖 1：在 Kaggle Notebook 中可以免費(fèi)使用 CPU、GPU 和 TPU。
如何準(zhǔn)備測(cè)試？
為了比較 CPU、GPU 和 TPU 完成常見(jiàn)數(shù)據(jù)科學(xué)任務(wù)時(shí)的性能，我們使用 tf_flowers 數(shù)據(jù)集訓(xùn)練了一個(gè)卷積神經(jīng)網(wǎng)絡(luò)，然后用三種不同的后端（GPU 是英偉達(dá) Tesla P100+2GHz、雙核的英特至強(qiáng) CPU，13 GB RAM；TPU 是谷歌 8 核的 TPUv3+2GHz、四核的英特爾至強(qiáng) CPU，16 GB RAM）將相同的代碼運(yùn)行了三遍。所附的 Kaggle 教程筆記（tutorial notebook）顯示，在一些情況下，TPU 性能最好。
例如：
用分片文檔數(shù)據(jù)集（例如 TFRecord）；
用 tf.data API 將訓(xùn)練數(shù)據(jù)傳到 TPU；
batch size 較大時(shí)（如 batch_size=128）。
在工作流程中加入這些步驟，可以避免常見(jiàn)的 I/O 瓶頸，否則這樣的 I/O 瓶頸會(huì)妨礙 TPU 發(fā)揮全部功能。用戶(hù)可以訪(fǎng)問(wèn) Kaggle TPU 官方文檔，了解其他一些在 TPU 上運(yùn)行代碼的優(yōu)化技巧。
官網(wǎng)鏈接：https://www.kaggle.com/docs/tpu
硬件性能如何
這三種硬件之間最明顯的區(qū)別在于使用 tf.keras 訓(xùn)練模型的速度。tf.keras 庫(kù)是最流行的機(jī)器學(xué)習(xí)框架之一，因?yàn)?tf.keras 可以更快、更方便地實(shí)驗(yàn)新想法。如果寫(xiě)代碼的時(shí)間短，則就可以撥出更多時(shí)間進(jìn)行計(jì)算。如果等代碼跑完的時(shí)間短，也就可以有更多時(shí)間評(píng)估新想法了。在機(jī)器學(xué)習(xí)比賽中，tf.keras 和 TPU 是一對(duì)很強(qiáng)大的組合！
圖 2：tf.keras 概貌。(A)tf.keras 是最流行的實(shí)現(xiàn)深度學(xué)習(xí)模型的框架之一；(B) 在設(shè)計(jì)深度學(xué)習(xí)模型時(shí)，tf.keras 可以快速進(jìn)行實(shí)驗(yàn)和迭代。
在下圖 3 第一個(gè)實(shí)驗(yàn)中，我們?cè)谌N硬件上運(yùn)行了相同的代碼（官方教程筆記中的改進(jìn)版），并且為了避免 CPU 和 GPU 出現(xiàn)內(nèi)存不足的問(wèn)題，batch size 也設(shè)置得很?。ㄔO(shè)為 16）。在這種情況下可以看出，訓(xùn)練 Xception 模型時(shí)，TPU 比 CPU 快了約 100 倍，比 GPU 快了約 3.5 倍，這是因?yàn)?TPU 處理批大小很大的數(shù)據(jù)時(shí)效率更高。我們也試著將 batch size 增加到 128，結(jié)果發(fā)現(xiàn) TPU 速度快了 2 倍，GPU 和 CPU 均出現(xiàn)了內(nèi)存不足的問(wèn)題。所以，與之前的實(shí)驗(yàn)相比，TPU 訓(xùn)練 Xception 模型的速度比 GPU 快 7 倍。
官方教程筆記：https://www.kaggle.com/mgornergoogle/five-flowers-with-keras-and-xception-on-tpu
圖 3：用 CPU、GPU 和 TPU 訓(xùn)練 Xception 模型（12 個(gè) epoch）所需時(shí)間對(duì)比。
如下圖 4 所示，模型訓(xùn)練的加速情況也與模型類(lèi)別有關(guān)，Xception 和 Vgg16 就比 ResNet50 表現(xiàn)更好。在這種邊界很大的情況下，模型訓(xùn)練速度是 TPU 唯一超過(guò) GPU 的地方。例如，當(dāng)預(yù)測(cè)少量樣本時(shí)，TPU 比 CPU 快了約 3 倍，但又比 GPU 慢了約 3 倍（在某些情況，比如在 batch size 很大的情況下做預(yù)測(cè)，TPU 表現(xiàn)出色，但本實(shí)驗(yàn)中不存在這樣的情況）。
圖 4：模型類(lèi)型和任務(wù)類(lèi)型不同，加速情況也不同。(A) 與 ResNet50 相比，Xception 和 Vgg16 的提速更為明顯。(B) 在諸如少量樣本預(yù)測(cè)等特定任務(wù)中，GPU 的表現(xiàn)比 TPU 好。
為了使結(jié)果更為充分，我們注意到 Yu Emma Wang 等人 [1] 開(kāi)發(fā)了一個(gè)叫做 ParaDnn 的嚴(yán)格基準(zhǔn)，這一基準(zhǔn)可以用來(lái)比較不同硬件訓(xùn)練機(jī)器學(xué)習(xí)模型時(shí)的性能。利用 Yu Emma Wang 等人 [1] 的方法，可以總結(jié)出：用 TPU 代替 GPU 時(shí)，參數(shù)化模型性能可以提升 1 到 10 倍，使用模型的性能也可以提升 3 到 6.8 倍（下圖 5）。分片數(shù)據(jù)、很大的數(shù)據(jù)批和大模型結(jié)合在一起時(shí)，TPU 是最佳選擇。
圖 5：Wang[1] 論文中的重要發(fā)現(xiàn)。當(dāng) batch size 和 CNN 模型都較大時(shí)，TPU 的性能最好。你可以在 Kumar[2] 和 Jouppi[3] 的論文中，找到其他基準(zhǔn)實(shí)驗(yàn)。
[1] https://arxiv.org/abs/1907.10701
訓(xùn)練模型時(shí)的價(jià)格考量
雖然平等地比較了硬件，但它們?cè)趦r(jià)格上有相當(dāng)大的差異。TPU（谷歌 TPU v3 每小時(shí) 8.00 美元，GCP 上可以按需選擇的 TPU v2 每小時(shí) 4.50 美元）比 GPU（英偉達(dá) Tesla P100 每小時(shí) 1.46 美元）貴了大概五倍。雖然如此，如果你想優(yōu)化成本，那還是應(yīng)該選擇 TPU，因?yàn)樵谟?xùn)練相同模型的情況下，TPU 的訓(xùn)練速度至少快了 5 倍。
當(dāng)數(shù)據(jù)以分片格式儲(chǔ)存在 GCS bucket，然后以大 batch size 傳遞到 TPU 時(shí)，模型訓(xùn)練會(huì)提速約 5 倍，因此建議熟悉 tf.data API 的用戶(hù)使用 TPU。
有些機(jī)器學(xué)習(xí)實(shí)踐者優(yōu)先考慮模型訓(xùn)練時(shí)間，而不是模型訓(xùn)練成本。因此，對(duì)于想盡快完成模型訓(xùn)練的人來(lái)說(shuō)，TPU 是最佳選擇。在訓(xùn)練模型上花更少的時(shí)間，就可以花更多時(shí)間考慮新想法。但請(qǐng)不要一味聽(tīng)信本文內(nèi)容——你可以通過(guò) Kaggle Notebooks 在 CPU、GPU 和 TPU 上免費(fèi)運(yùn)行代碼并評(píng)估結(jié)果。Kaggle 用戶(hù)們已經(jīng)成功在 TPU 上運(yùn)行了文本數(shù)據(jù)，并從中獲得了許多樂(lè)趣。
下面這篇文章的作者描述了是如何用 TPU 訓(xùn)練 BERT 模型，并在最近的 Kaggle 競(jìng)賽中贏得了 8000 美元獎(jiǎng)金（二等獎(jiǎng)）的。
文章鏈接：https://www.kaggle.com/c/tensorflow2-question-answering/discussion/127333
該選哪種硬件呢？
總之，如果有通用性和大內(nèi)存需求的話(huà)，我們建議你使用 CPU。當(dāng)你想加快各種數(shù)據(jù)科學(xué)流程時(shí)，GPU 是很好的替代品。如果想要盡可能快地完成模型訓(xùn)練的話(huà)，TPU 則是最佳選擇。
針對(duì)所用硬件來(lái)優(yōu)化代碼，可以得到更好的結(jié)果。我們認(rèn)為，對(duì)分別針對(duì) GPU 代碼和 TPU 代碼的運(yùn)行時(shí)（runtime）進(jìn)行比較也很有意思。例如，用像 RAPIDS.ai 這樣的 GPU 加速庫(kù)訓(xùn)練梯度提升模型，再用像 tf.keras 這樣的 TPU 加速庫(kù)訓(xùn)練深度學(xué)習(xí)模型，比較二者的訓(xùn)練時(shí)間，這也是很有意思的。
那么，訓(xùn)練準(zhǔn)確的機(jī)器學(xué)習(xí)模型最少需要多長(zhǎng)時(shí)間？一天內(nèi)可以評(píng)價(jià)多少不同的想法？結(jié)合 tf.keras，TPU 讓機(jī)器學(xué)習(xí)從業(yè)人員寫(xiě)代碼的時(shí)間變短了，等代碼跑完的時(shí)間也變短了，這樣就可以留出更多時(shí)間評(píng)估新想法，并在 Kaggle 大賽中提升自己作品的性能。
原文鏈接：https://towardsdatascience.com/when-to-use-cpus-vs-gpus-vs-tpus-in-a-kaggle-competition-9af708a8c3eb
本文為機(jī)器之心編譯，轉(zhuǎn)載請(qǐng)聯(lián)系本公眾號(hào)獲得授權(quán)。
