岳排槐 發(fā)自 凹非寺
量子位 出品 | 公眾號(hào) QbitAI

對(duì)于大多數(shù)搞深度學(xué)習(xí)的人來說，英偉達(dá)GPU之外其實(shí)沒有更多選擇。然而谷歌TPU芯片的出現(xiàn)，有望改變這一現(xiàn)狀。

這個(gè)大名鼎鼎的AI芯片，即是谷歌各種AI應(yīng)用和服務(wù)背后的支撐，也是名噪天下的AlphaGo背后的基礎(chǔ)。碾壓人類頂級(jí)圍棋選手，只需要四塊TPU。

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

在首次公布9個(gè)月后，谷歌TPU終于面向大眾開放。10天前，谷歌的Cloud TPU正式發(fā)布。只需要每小時(shí)6.5美元，你也有可能用上谷歌TPU。

到底Cloud TPU實(shí)力如何？RiseML（riseml.com）最近做了一次對(duì)比評(píng)測(cè)。

云上的TPU

首先簡(jiǎn)單介紹一下測(cè)試對(duì)象。

第一代TPU面向推理，而第二代的重點(diǎn)是加速訓(xùn)練。在TPUv2的核心里，一個(gè)脈動(dòng)陣列（Systolic array）負(fù)責(zé)執(zhí)行矩陣乘法，這在深度學(xué)習(xí)中被大量使用。

根據(jù)Jeff Dean此前放出的PPT顯示，每個(gè)Cloud TPU由四個(gè)TPUv2芯片組成。每個(gè)芯片有16GB內(nèi)存和兩個(gè)內(nèi)核，每個(gè)內(nèi)核有兩個(gè)矩陣乘法單元。

兩個(gè)內(nèi)核能夠提供45TFLOPs算力，所以每個(gè)Cloud TPU能提供180TFLOPs算力和64GB內(nèi)存。作為對(duì)比，這一代英偉達(dá)V100 GPU提供125TFLOPs算力和16GB內(nèi)存。

當(dāng)你獲得配額后，就能在谷歌云上啟動(dòng)Cloud TPU。無需（也沒有辦法）把一個(gè)Cloud TPU分配給指定的VM實(shí)例。每個(gè)Cloud TPU會(huì)有一個(gè)名字和IP地址，供用戶提供TensorFlow代碼。

△ 創(chuàng)建一個(gè)新的Cloud TPU

Cloud TPU僅支持TensorFlow 1.6版本。除此之外，你的VM實(shí)例上不需要任何驅(qū)動(dòng)程序，因?yàn)榕cTPU進(jìn)行通信所需的所有代碼都由TensorFlow本身提供。在TPU上執(zhí)行的代碼經(jīng)過優(yōu)化，并由XLA進(jìn)行實(shí)時(shí)編譯，XLA也是TensorFlow的一部分。

為了有效使用TPU，你的代碼應(yīng)該建立在高級(jí)Estimator抽象上。然后可以用TPUEstimator^[1]來執(zhí)行很多必要的任務(wù)，這也能更有效的利用TPU。例如，這可以為TPU設(shè)置數(shù)據(jù)隊(duì)列并在不同的核心之間并行計(jì)算。

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

一旦完成所有設(shè)置，就能像普通情況一樣運(yùn)行TensorFlow代碼。TPU將在啟動(dòng)過程中被發(fā)現(xiàn)，然后計(jì)算圖被編譯并傳輸?shù)絋PU。有意思的是，TPU還可以直接從云存儲(chǔ)中讀取和寫入，存儲(chǔ)檢查點(diǎn)或者事件摘要。當(dāng)然你需要提供相應(yīng)的寫入和訪問權(quán)限。

評(píng)測(cè)設(shè)置

這個(gè)評(píng)測(cè)最想得到的回答，當(dāng)然是TPU速度有多快。

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

接下來的評(píng)測(cè)，基于ResNet和Inception。

我們還想知道，沒有針對(duì)TPU進(jìn)行優(yōu)化的模型運(yùn)行起來什么樣，所以還有一個(gè)進(jìn)行文本分類的LSTM模型運(yùn)行在TPU上。而且這還是一個(gè)小模型，我們也想看看效果，因?yàn)橥ǔ９雀杞ㄗh在TPU上運(yùn)行較大的模型。

所有的模型，都有會(huì)分別在單個(gè)Cloud TPU和單個(gè)英偉達(dá)P100、V100 GPU上進(jìn)行訓(xùn)練，然后進(jìn)行速度比較。當(dāng)然，徹底的比較還應(yīng)包括模型的最終質(zhì)量、收斂性等。但是這次的評(píng)測(cè)，只關(guān)注了訓(xùn)練速度，更多情況稍后再詳細(xì)研究。

在TPU和P100上的實(shí)驗(yàn)，運(yùn)行于谷歌云平臺(tái)的n1-standard-16實(shí)例（16 vCPUs Intel Haswell, 60 GB memory）。在V100上的實(shí)驗(yàn)，使用了亞馬遜云的p3.2xlarge實(shí)例（8 vCPUs, 60 GB memory）。

所有的系統(tǒng)都運(yùn)行于Ubuntu 16.04。對(duì)于TPU，我們從PyPi倉庫安裝了TensorFlow 1.6.0-rc1。GPU實(shí)驗(yàn)中運(yùn)行了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

實(shí)驗(yàn)結(jié)果：TPU優(yōu)化的模型

首先來看看第一組結(jié)果，針對(duì)TPU優(yōu)化過的模型表現(xiàn)如何。下面，可以看到性能對(duì)比，衡量標(biāo)準(zhǔn)是每秒能處理的圖片數(shù)。

TPU的batch大小是1024，GPU是128。對(duì)于GPU，我們使用了TensorFlow基準(zhǔn)倉庫的實(shí)現(xiàn)，地址：https://github.com/tensorflow/benchmarks。訓(xùn)練數(shù)據(jù)是谷歌提供的偽ImageNet數(shù)據(jù)集，存儲(chǔ)在云端（為TPU）和本地磁盤（為GPU）。

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

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

實(shí)驗(yàn)結(jié)果：自定義的LSTM模型

我們自定義的模型時(shí)雙向LSTM，使用1024隱藏單元進(jìn)行文本分類。LSTM是當(dāng)今NLP的一個(gè)基本構(gòu)建模塊，而上述官方模型都是基于計(jì)算機(jī)視覺。

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

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

由于這個(gè)模型相對(duì)較小，而且沒有任何優(yōu)化調(diào)整，所以結(jié)果喜人啊。當(dāng)然bug還沒修復(fù)，所以上述結(jié)果只是初步結(jié)果，僅供參考。

結(jié)論

上述參與測(cè)試的結(jié)果表明，谷歌Cloud TPU與英偉達(dá)最新一代GPU相比，性能更好而且經(jīng)濟(jì)實(shí)惠。雖然谷歌為TPU進(jìn)行了更大規(guī)模模型的優(yōu)化，但小型模型仍然受益于此。總的來說，盡管只是beta測(cè)試階段，但Cloud TPU表現(xiàn)已經(jīng)很好了。

RiseML最后給出結(jié)論：

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

如何申請(qǐng)使用

最后說說，怎么才能用上數(shù)量有限的Cloud TPU。

要使用beta版的Cloud TPU，需要填個(gè)表，描述一下你要用TPU干什么，向谷歌申請(qǐng)配額：https://services.google.com/fb/forms/cloud-tpu-beta-request/

谷歌說，會(huì)盡快讓你用上Cloud TPU。

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

一家是投資公司Two Sigma。他們的深度學(xué)習(xí)研究現(xiàn)在主要在云上進(jìn)行，該公司CTO Alfred Spector說：“將TensorFlow工作負(fù)載轉(zhuǎn)移到TPU上，大大降低了編程新模型的復(fù)雜性，縮短了訓(xùn)練時(shí)間?！?/p>

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