僅需1/5成本：TPU是如何超越GPU，成為深度學習首選處理器的

很多讀者可能分不清楚 CPU、GPU 和 TPU 之間的區別，因此 Google Cloud 將在這篇部落格中簡要介紹它們之間的區別，並討論為什麼 TPU 能加速深度學習。

張量處理單元（TPU）是一種定製化的 ASIC 晶片，它由谷歌從頭設計，並專門用於機器學習工作負載。TPU 為谷歌的主要產品提供了計算支援，包括翻譯、照片、搜尋助理和 Gmail 等。Cloud TPU 將 TPU 作為可擴充套件的雲端計算資源，併為所有在 Google Cloud 上執行尖端 ML 模型的開發者與資料科學家提供計算資源。在 Google Next’18 中，我們宣佈 TPU v2 現在已經得到使用者的廣泛使用，包括那些免費試用使用者，而 TPU v3 目前已經發布了內部測試版。

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

神經網路如何運算

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

如果影象為 28×28 畫素的灰度圖，那麼它可以轉化為包含 784 個元素的向量。神經元會接收所有 784 個值，並將它們與引數值（上圖紅線）相乘，因此才能識別為「8」。其中引數值的作用類似於用「濾波器」從資料中抽取特徵，因而能計算輸入影象與「8」之間的相似性：

這是對神經網路做資料分類最基礎的解釋，即將資料與對應的引數相乘（上圖兩種顏色的點），並將它們加在一起（上圖右側收集計算結果）。如果我們能得到最高的預測值，那麼我們會發現輸入資料與對應引數非常匹配，這也就最可能是正確的答案。

簡單而言，神經網路在資料和引數之間需要執行大量的乘法和加法。我們通常會將這些乘法與加法組合為矩陣運算，這在我們大學的線性代數中會提到。所以關鍵點是我們該如何快速執行大型矩陣運算，同時還需要更小的能耗。

CPU 如何執行

因此 CPU 如何來執行這樣的大型矩陣運算任務呢？一般 CPU 是基於馮諾依曼架構的通用處理器，這意味著 CPU 與軟體和記憶體的執行方式如下：

CPU 最大的優勢是靈活性。通過馮諾依曼架構，我們可以為數百萬的不同應用載入任何軟體。我們可以使用 CPU 處理文字、控制火箭引擎、執行銀行交易或者使用神經網路分類影象。

但是，由於 CPU 非常靈活，硬體無法一直了解下一個計算是什麼，直到它讀取了軟體的下一個指令。CPU 必須在內部將每次計算的結果儲存到記憶體中（也被稱為暫存器或 L1 快取）。記憶體訪問成為 CPU 架構的不足，被稱為馮諾依曼瓶頸。雖然神經網路的大規模運算中的每一步都是完全可預測的，每一個 CPU 的算術邏輯單元（ALU，控制乘法器和加法器的元件）都只能一個接一個地執行它們，每一次都需要訪問記憶體，限制了總體吞吐量，並需要大量的能耗。

GPU 如何工作

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

這種 GPU 架構在有大量並行化的應用中工作得很好，例如在神經網路中的矩陣乘法。實際上，相比 CPU，GPU 在深度學習的典型訓練工作負載中能實現高几個數量級的吞吐量。這正是為什麼 GPU 是深度學習中最受歡迎的處理器架構。

但是，GPU 仍然是一種通用的處理器，必須支援幾百萬種不同的應用和軟體。這又把我們帶回到了基礎的問題，馮諾依曼瓶頸。在每次幾千個 ALU 的計算中，GPU 都需要訪問暫存器或共享記憶體來讀取和儲存中間計算結果。因為 GPU 在其 ALU 上執行更多的平行計算，它也會成比例地耗費更多的能量來訪問記憶體，同時也因為複雜的線路而增加 GPU 的物理空間佔用。

TPU 如何工作

當谷歌設計 TPU 的時候，我們構建了一種領域特定的架構。這意味著，我們沒有設計一種通用的處理器，而是專用於神經網路工作負載的矩陣處理器。TPU 不能執行文字處理軟體、控制火箭引擎或執行銀行業務，但它們可以為神經網路處理大量的乘法和加法運算，同時 TPU 的速度非常快、能耗非常小且物理空間佔用也更小。

其主要助因是對馮諾依曼瓶頸的大幅度簡化。因為該處理器的主要任務是矩陣處理，TPU 的硬體設計者知道該運算過程的每個步驟。因此他們放置了成千上萬的乘法器和加法器並將它們直接連線起來，以構建那些運算子的物理矩陣。這被稱作脈動陣列（Systolic Array）架構。在 Cloud TPU v2 的例子中，有兩個 128X128 的脈動陣列，在單個處理器中整合了 32768 個 ALU 的 16 位浮點值。

我們來看看一個脈動陣列如何執行神經網路計算。首先，TPU 從記憶體載入引數到乘法器和加法器的矩陣中。

然後，TPU 從記憶體載入資料。當每個乘法被執行後，其結果將被傳遞到下一個乘法器，同時執行加法。因此結果將是所有資料和引數乘積的和。在大量計算和資料傳遞的整個過程中，不需要執行任何的記憶體訪問。

這就是為什麼 TPU 可以在神經網路運算上達到高計算吞吐量，同時能耗和物理空間都很小。

好處：成本降低至 1/5

因此使用 TPU 架構的好處就是：降低成本。以下是截至 2018 年 8 月（寫這篇文章的時候）Cloud TPU v2 的使用價格。

史丹佛大學釋出了深度學習和推理的基準套裝 DAWNBench。你可以在上面找到不同的任務、模型、計算平臺以及各自的基準結果的組合。

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

在 DAWNBench 比賽於 2018 年 4 月結束的時候，非 TPU 處理器的最低訓練成本是 72.40 美元（使用現場例項訓練 ResNet-50 達到 93% 準確率）。而使用 Cloud TPU v2 搶佔式計價，你可以在 12.87 美元的價格完成相同的訓練結果。這僅相當於非 TPU 的不到 1/5 的成本。這正是神經網路領域特定架構的威力之所在。