張量處理單元（TPU）是一種定制化的 ASIC 芯片，它由谷歌從頭設(shè)計(jì)，并專(zhuān)門(mén)用于機(jī)器學(xué)習(xí)工作負(fù)載。TPU 為谷歌的主要產(chǎn)品提供了計(jì)算支持，包括翻譯、照片、搜索助理和 Gmail 等。

Cloud TPU 將 TPU 作為可擴(kuò)展的云計(jì)算資源，并為所有在 Google Cloud 上運(yùn)行尖端 ML 模型的開(kāi)發(fā)者與數(shù)據(jù)科學(xué)家提供計(jì)算資源。在 Google Next’18 中，我們宣布 TPU v2 現(xiàn)在已經(jīng)得到用戶(hù)的廣泛使用，包括那些免費(fèi)試用用戶(hù)，而 TPU v3 目前已經(jīng)發(fā)布了內(nèi)部測(cè)試版。

第三代 Cloud TPU

如上為 tpudemo.com 截圖，該網(wǎng)站 PPT 解釋了 TPU 的特性與定義。在本文中，我們將關(guān)注 TPU 某些特定的屬性。

神經(jīng)網(wǎng)絡(luò)如何運(yùn)算

在我們對(duì)比 CPU、GPU 和 TPU 之前，我們可以先了解到底機(jī)器學(xué)習(xí)或神經(jīng)網(wǎng)絡(luò)需要什么樣的計(jì)算。如下所示，假設(shè)我們使用單層神經(jīng)網(wǎng)絡(luò)識(shí)別手寫(xiě)數(shù)字。

如果圖像為 28×28 像素的灰度圖，那么它可以轉(zhuǎn)化為包含 784 個(gè)元素的向量。神經(jīng)元會(huì)接收所有 784 個(gè)值，并將它們與參數(shù)值（上圖紅線）相乘，因此才能識(shí)別為「8」。其中參數(shù)值的作用類(lèi)似于用「濾波器」從數(shù)據(jù)中抽取特征，因而能計(jì)算輸入圖像與「8」之間的相似性：

這是對(duì)神經(jīng)網(wǎng)絡(luò)做數(shù)據(jù)分類(lèi)最基礎(chǔ)的解釋?zhuān)磳?shù)據(jù)與對(duì)應(yīng)的參數(shù)相乘（上圖兩種顏色的點(diǎn)），并將它們加在一起（上圖右側(cè)收集計(jì)算結(jié)果）。如果我們能得到最高的預(yù)測(cè)值，那么我們會(huì)發(fā)現(xiàn)輸入數(shù)據(jù)與對(duì)應(yīng)參數(shù)非常匹配，這也就最可能是正確的答案。

簡(jiǎn)單而言，神經(jīng)網(wǎng)絡(luò)在數(shù)據(jù)和參數(shù)之間需要執(zhí)行大量的乘法和加法。我們通常會(huì)將這些乘法與加法組合為矩陣運(yùn)算，這在我們大學(xué)的線性代數(shù)中會(huì)提到。所以關(guān)鍵點(diǎn)是我們?cè)撊绾慰焖賵?zhí)行大型矩陣運(yùn)算，同時(shí)還需要更小的能耗。

CPU 如何運(yùn)行

因此 CPU 如何來(lái)執(zhí)行這樣的大型矩陣運(yùn)算任務(wù)呢？一般 CPU 是基于馮諾依曼架構(gòu)的通用處理器，這意味著 CPU 與軟件和內(nèi)存的運(yùn)行方式如下：

CPU 如何運(yùn)行：該動(dòng)圖僅展示了概念性原理，并不反映 CPU 的實(shí)際運(yùn)算行為。

CPU 最大的優(yōu)勢(shì)是靈活性。通過(guò)馮諾依曼架構(gòu)，我們可以為數(shù)百萬(wàn)的不同應(yīng)用加載任何軟件。我們可以使用 CPU 處理文字、控制火箭引擎、執(zhí)行銀行交易或者使用神經(jīng)網(wǎng)絡(luò)分類(lèi)圖像。

但是，由于 CPU 非常靈活，硬件無(wú)法一直了解下一個(gè)計(jì)算是什么，直到它讀取了軟件的下一個(gè)指令。CPU 必須在內(nèi)部將每次計(jì)算的結(jié)果保存到內(nèi)存中（也被稱(chēng)為寄存器或 L1 緩存）。內(nèi)存訪問(wèn)成為 CPU 架構(gòu)的不足，被稱(chēng)為馮諾依曼瓶頸。雖然神經(jīng)網(wǎng)絡(luò)的大規(guī)模運(yùn)算中的每一步都是完全可預(yù)測(cè)的，每一個(gè) CPU 的算術(shù)邏輯單元（ALU，控制乘法器和加法器的組件）都只能一個(gè)接一個(gè)地執(zhí)行它們，每一次都需要訪問(wèn)內(nèi)存，限制了總體吞吐量，并需要大量的能耗。

GPU 如何工作

為了獲得比 CPU 更高的吞吐量，GPU 使用一種簡(jiǎn)單的策略：在單個(gè)處理器中使用成千上萬(wàn)個(gè) ALU?，F(xiàn)代 GPU 通常在單個(gè)處理器中擁有 2500-5000 個(gè) ALU，意味著你可以同時(shí)執(zhí)行數(shù)千次乘法和加法運(yùn)算。

GPU 如何工作：這個(gè)動(dòng)畫(huà)僅用于概念展示。并不反映真實(shí)處理器的實(shí)際工作方式。

這種 GPU 架構(gòu)在有大量并行化的應(yīng)用中工作得很好，例如在神經(jīng)網(wǎng)絡(luò)中的矩陣乘法。實(shí)際上，相比 CPU，GPU 在深度學(xué)習(xí)的典型訓(xùn)練工作負(fù)載中能實(shí)現(xiàn)高幾個(gè)數(shù)量級(jí)的吞吐量。這正是為什么 GPU 是深度學(xué)習(xí)中最受歡迎的處理器架構(gòu)。

但是，GPU 仍然是一種通用的處理器，必須支持幾百萬(wàn)種不同的應(yīng)用和軟件。這又把我們帶回到了基礎(chǔ)的問(wèn)題，馮諾依曼瓶頸。在每次幾千個(gè) ALU 的計(jì)算中，GPU 都需要訪問(wèn)寄存器或共享內(nèi)存來(lái)讀取和保存中間計(jì)算結(jié)果。因?yàn)?GPU 在其 ALU 上執(zhí)行更多的并行計(jì)算，它也會(huì)成比例地耗費(fèi)更多的能量來(lái)訪問(wèn)內(nèi)存，同時(shí)也因?yàn)閺?fù)雜的線路而增加 GPU 的物理空間占用。

TPU 如何工作

當(dāng)谷歌設(shè)計(jì) TPU 的時(shí)候，我們構(gòu)建了一種領(lǐng)域特定的架構(gòu)。這意味著，我們沒(méi)有設(shè)計(jì)一種通用的處理器，而是專(zhuān)用于神經(jīng)網(wǎng)絡(luò)工作負(fù)載的矩陣處理器。TPU 不能運(yùn)行文本處理軟件、控制火箭引擎或執(zhí)行銀行業(yè)務(wù)，但它們可以為神經(jīng)網(wǎng)絡(luò)處理大量的乘法和加法運(yùn)算，同時(shí) TPU 的速度非常快、能耗非常小且物理空間占用也更小。

其主要助因是對(duì)馮諾依曼瓶頸的大幅度簡(jiǎn)化。因?yàn)樵撎幚砥鞯闹饕蝿?wù)是矩陣處理，TPU 的硬件設(shè)計(jì)者知道該運(yùn)算過(guò)程的每個(gè)步驟。因此他們放置了成千上萬(wàn)的乘法器和加法器并將它們直接連接起來(lái)，以構(gòu)建那些運(yùn)算符的物理矩陣。這被稱(chēng)作脈動(dòng)陣列（Systolic Array）架構(gòu)。在 Cloud TPU v2 的例子中，有兩個(gè) 128X128 的脈動(dòng)陣列，在單個(gè)處理器中集成了 32768 個(gè) ALU 的 16 位浮點(diǎn)值。

我們來(lái)看看一個(gè)脈動(dòng)陣列如何執(zhí)行神經(jīng)網(wǎng)絡(luò)計(jì)算。首先，TPU 從內(nèi)存加載參數(shù)到乘法器和加法器的矩陣中。

然后，TPU 從內(nèi)存加載數(shù)據(jù)。當(dāng)每個(gè)乘法被執(zhí)行后，其結(jié)果將被傳遞到下一個(gè)乘法器，同時(shí)執(zhí)行加法。因此結(jié)果將是所有數(shù)據(jù)和參數(shù)乘積的和。在大量計(jì)算和數(shù)據(jù)傳遞的整個(gè)過(guò)程中，不需要執(zhí)行任何的內(nèi)存訪問(wèn)。

這就是為什么 TPU 可以在神經(jīng)網(wǎng)絡(luò)運(yùn)算上達(dá)到高計(jì)算吞吐量，同時(shí)能耗和物理空間都很小。

好處：成本降低至 1/5

因此使用 TPU 架構(gòu)的好處就是：降低成本。以下是截至 2018 年 8 月（寫(xiě)這篇文章的時(shí)候）Cloud TPU v2 的使用價(jià)格。

Cloud TPU v2 的價(jià)格，截至 2018 年 8 月。

斯坦福大學(xué)發(fā)布了深度學(xué)習(xí)和推理的基準(zhǔn)套裝 DAWNBench。你可以在上面找到不同的任務(wù)、模型、計(jì)算平臺(tái)以及各自的基準(zhǔn)結(jié)果的組合。

DAWNBench：https://dawn.cs.stanford.edu/benchmark/

在 DAWNBench 比賽于 2018 年 4 月結(jié)束的時(shí)候，非 TPU 處理器的最低訓(xùn)練成本是 72.40 美元（使用現(xiàn)場(chǎng)實(shí)例訓(xùn)練 ResNet-50 達(dá)到 93% 準(zhǔn)確率）。而使用 Cloud TPU v2 搶占式計(jì)價(jià)，你可以在 12.87 美元的價(jià)格完成相同的訓(xùn)練結(jié)果。這僅相當(dāng)于非 TPU 的不到 1/5 的成本。這正是神經(jīng)網(wǎng)絡(luò)領(lǐng)域特定架構(gòu)的威力之所在。