谷歌這個大殺器要讓英偉達慌了，實戰評測：TPU相比GPU簡直又快又省

岳排槐 發自 凹非寺

量子位 出品 | 公眾號 QbitAI

對於大多數搞深度學習的人來說，英偉達GPU之外其實沒有更多選擇。然而谷歌TPU晶元的出現，有望改變這一現狀。

這個大名鼎鼎的AI晶元，即是谷歌各種AI應用和服務背後的支撐，也是名噪天下的AlphaGo背後的基礎。碾壓人類頂級圍棋選手，只需要四塊TPU。

但此前，很少有「外人」能一嘗TPU的滋味。

在首次公布9個月後，谷歌TPU終於面向大眾開放。10天前，谷歌的Cloud TPU正式發布。只需要每小時6.5美元，你也有可能用上谷歌TPU。

到底Cloud TPU實力如何？RiseML（riseml.com）最近做了一次對比評測。

雲上的TPU

首先簡單介紹一下測試對象。

第一代TPU面向推理，而第二代的重點是加速訓練。在TPUv2的核心裡，一個脈動陣列（Systolic array）負責執行矩陣乘法，這在深度學習中被大量使用。

根據Jeff Dean此前放出的PPT顯示，每個Cloud TPU由四個TPUv2晶元組成。每個晶元有16GB內存和兩個內核，每個內核有兩個矩陣乘法單元。

兩個內核能夠提供45TFLOPs算力，所以每個Cloud TPU能提供180TFLOPs算力和64GB內存。作為對比，這一代英偉達V100 GPU提供125TFLOPs算力和16GB內存。

當你獲得配額後，就能在谷歌雲上啟動Cloud TPU。無需（也沒有辦法）把一個Cloud TPU分配給指定的VM實例。每個Cloud TPU會有一個名字和IP地址，供用戶提供TensorFlow代碼。

創建一個新的Cloud TPU

Cloud TPU僅支持TensorFlow 1.6版本。除此之外，你的VM實例上不需要任何驅動程序，因為與TPU進行通信所需的所有代碼都由TensorFlow本身提供。在TPU上執行的代碼經過優化，並由XLA進行實時編譯，XLA也是TensorFlow的一部分。

為了有效使用TPU，你的代碼應該建立在高級Estimator抽象上。然後可以用TPUEstimator[1]來執行很多必要的任務，這也能更有效的利用TPU。例如，這可以為TPU設置數據隊列並在不同的核心之間並行計算。

[1] https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUEstimator

一旦完成所有設置，就能像普通情況一樣運行TensorFlow代碼。TPU將在啟動過程中被發現，然後計算圖被編譯並傳輸到TPU。有意思的是，TPU還可以直接從雲存儲中讀取和寫入，存儲檢查點或者事件摘要。當然你需要提供相應的寫入和訪問許可權。

評測設置

這個評測最想得到的回答，當然是TPU速度有多快。

TensorFlow在GitHub上提供了一個針對TPU的模型倉庫。地址在：https://github.com/tensorflow/tpu。

接下來的評測，基於ResNet和Inception。

我們還想知道，沒有針對TPU進行優化的模型運行起來什麼樣，所以還有一個進行文本分類的LSTM模型運行在TPU上。而且這還是一個小模型，我們也想看看效果，因為通常谷歌建議在TPU上運行較大的模型。

所有的模型，都有會分別在單個Cloud TPU和單個英偉達P100、V100 GPU上進行訓練，然後進行速度比較。當然，徹底的比較還應包括模型的最終質量、收斂性等。但是這次的評測，只關注了訓練速度，更多情況稍後再詳細研究。

在TPU和P100上的實驗，運行於谷歌雲平台的n1-standard-16實例（16 vCPUs Intel Haswell, 60 GB memory）。在V100上的實驗，使用了亞馬遜雲的p3.2xlarge實例（8 vCPUs, 60 GB memory）。

所有的系統都運行於Ubuntu 16.04。對於TPU，我們從PyPi倉庫安裝了TensorFlow 1.6.0-rc1。GPU實驗中運行了nvidia-docker[2]，使用了TensorFlow 1.5（tensorflow:1.5.0-gpu-py3）其中包括CUDA 9.0和cuDNN 7.0。

[2] https://github.com/NVIDIA/nvidia-docker

實驗結果：TPU優化的模型

首先來看看第一組結果，針對TPU優化過的模型表現如何。下面，可以看到性能對比，衡量標準是每秒能處理的圖片數。

TPU的batch大小是1024，GPU是128。對於GPU，我們使用了TensorFlow基準倉庫的實現，地址：https://github.com/tensorflow/benchmarks。訓練數據是谷歌提供的偽ImageNet數據集，存儲在雲端（為TPU）和本地磁碟（為GPU）。

在ResNet-50上，單個Cloud TPU比單個P100快8.4倍，比V100快5.1倍。對於InceptionV3，結果差不多，分別快8.4倍和4.8倍。另外，如果把精度降低（fp16），V100的提速更加明顯。

除了速度之外，成本也是重要考慮因素。下面這個表格顯示，如果都在雲端計算，TPU的性價比還是最高的。

實驗結果：自定義的LSTM模型

我們自定義的模型時雙向LSTM，使用1024隱藏單元進行文本分類。LSTM是當今NLP的一個基本構建模塊，而上述官方模型都是基於計算機視覺。

源代碼已經使用了Estimator框架，所以將其用於TPUEstimator非常簡單。但是需要注意的是：在TPU上我們無法讓模型收斂，而GPU上相同模型（batch大小等）工作正常。這應該是一個bug，或者是代碼問題或是是TensorFlow的問題。

實驗結果表明，在自定義的LSTM模型上，TPU還是更快。TPU（21402 examples/s）比P100（1658 examples/s）快16.9倍，比V100（2778 examples/s）快7.7倍。

由於這個模型相對較小，而且沒有任何優化調整，所以結果喜人啊。當然bug還沒修復，所以上述結果只是初步結果，僅供參考。

結論

上述參與測試的結果表明，谷歌Cloud TPU與英偉達最新一代GPU相比，性能更好而且經濟實惠。雖然谷歌為TPU進行了更大規模模型的優化，但小型模型仍然受益於此。總的來說，儘管只是beta測試階段，但Cloud TPU表現已經很好了。

RiseML最後給出結論：

一旦TPU能夠容納更多的用戶使用，就可以成為英偉達GPU真正的替代者。

如何申請使用

最後說說，怎麼才能用上數量有限的Cloud TPU。

要使用beta版的Cloud TPU，需要填個表，描述一下你要用TPU幹什麼，向谷歌申請配額：https://services.google.com/fb/forms/cloud-tpu-beta-request/

谷歌說，會儘快讓你用上Cloud TPU。

此前的博客文章中，谷歌提到了兩家客戶使用Cloud TPU的感受。

一家是投資公司Two Sigma。他們的深度學習研究現在主要在雲上進行，該公司CTO Alfred Spector說：「將TensorFlow工作負載轉移到TPU上，大大降低了編程新模型的複雜性，縮短了訓練時間。」

另一家是共享出行公司Lyft。深度學習正在成為這家公司無人車研究的重要組成部分。

—完—

加入社群

量子位AI社群13群開始招募啦，歡迎對AI感興趣的同學，加小助手微信qbitbot5入群；

此外，量子位專業細分群(自動駕駛、CV、NLP、機器學習等)正在招募，面向正在從事相關領域的工程師及研究人員。

進群請加小助手微信號qbitbot5，並務必備註相應群的關鍵詞~通過審核後我們將邀請進群。（專業群審核較嚴，敬請諒解）

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。

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