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

在本文中，我們將關注 TPU 某些特定的屬性。

神經(jīng)網(wǎng)絡如何運算

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

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

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

簡單而言，神經(jīng)網(wǎng)絡在數(shù)據(jù)和參數(shù)之間需要執(zhí)行大量的乘法和加法。我們通常會將這些乘法與加法組合為矩陣運算，這在我們大學的線性代數(shù)中會提到。所以關鍵點是我們該如何快速執(zhí)行大型矩陣運算，同時還需要更小的能耗。

CPU 如何運行

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

圖：CPU 如何運行

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

但是，由于 CPU 非常靈活，硬件無法一直了解下一個計算是什么，直到它讀取了軟件的下一個指令。CPU 必須在內(nèi)部將每次計算的結果保存到內(nèi)存中（也被稱為寄存器或 L1 緩存）。內(nèi)存訪問成為 CPU 架構的不足，被稱為馮諾依曼瓶頸。

雖然神經(jīng)網(wǎng)絡的大規(guī)模運算中的每一步都是完全可預測的，每一個 CPU 的算術邏輯單元（ALU，控制乘法器和加法器的組件）都只能一個接一個地執(zhí)行它們，每一次都需要訪問內(nèi)存，限制了總體吞吐量，并需要大量的能耗。

GPU 如何工作

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

圖：GPU 如何工作

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

但是，GPU 仍然是一種通用的處理器，必須支持幾百萬種不同的應用和軟件。這又把我們帶回到了基礎的問題，馮諾依曼瓶頸。在每次幾千個 ALU 的計算中，GPU 都需要訪問寄存器或共享內(nèi)存來讀取和保存中間計算結果。

因為 GPU 在其 ALU 上執(zhí)行更多的并行計算，它也會成比例地耗費更多的能量來訪問內(nèi)存，同時也因為復雜的線路而增加 GPU 的物理空間占用。

TPU 如何工作

當谷歌設計 TPU 的時候，我們構建了一種領域特定的架構。這意味著，我們沒有設計一種通用的處理器，而是專用于神經(jīng)網(wǎng)絡工作負載的矩陣處理器。

TPU 不能運行文本處理軟件、控制火箭引擎或執(zhí)行銀行業(yè)務，但它們可以為神經(jīng)網(wǎng)絡處理大量的乘法和加法運算，同時 TPU 的速度非?？?、能耗非常小且物理空間占用也更小。

其主要助因是對馮諾依曼瓶頸的大幅度簡化。因為該處理器的主要任務是矩陣處理，TPU 的硬件設計者知道該運算過程的每個步驟。因此他們放置了成千上萬的乘法器和加法器并將它們直接連接起來，以構建那些運算符的物理矩陣。

這被稱作脈動陣列（Systolic Array）架構。在 Cloud TPU v2 的例子中，有兩個 128X128 的脈動陣列，在單個處理器中集成了 32768 個 ALU 的 16 位浮點值。

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

圖：TPU 如何工作

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

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

因此使用 TPU 架構的好處就是：成本降低至 1/5。