谷歌最便宜TPU值不值得買？TPU在執(zhí)行神經(jīng)網(wǎng)絡(luò)計(jì)算方面的優(yōu)勢

谷歌本月推出千元級搭載Edge TPU芯片的開發(fā)板，性能令人期待。本文以可視化圖形的方式，對比TPU、GPU和CPU，解釋了TPU在執(zhí)行神經(jīng)網(wǎng)絡(luò)計(jì)算方面的優(yōu)勢。

谷歌最便宜 TPU 值不值得買？

谷歌 Edge TPU 在本月初終于公布價(jià)格 —— 不足 1000 元人民幣，遠(yuǎn)低于 TPU。

實(shí)際上，Edge TPU 基本上就是機(jī)器學(xué)習(xí)的樹莓派，它是一個(gè)用 TPU 在邊緣進(jìn)行推理的設(shè)備。

Edge TPU(安裝在 Coral 開發(fā)板上）

云 vs 邊緣

Edge TPU顯然是在邊緣(edge)運(yùn)行的，但邊緣是什么呢？為什么我們不選擇在云上運(yùn)行所有東西呢？

在云中運(yùn)行代碼意味著你使用的CPU、GPU和TPU都是通過瀏覽器提供的。在云中運(yùn)行代碼的主要優(yōu)點(diǎn)是，你可以為特定的代碼分配必要的計(jì)算能力（訓(xùn)練大型模型可能需要大量的計(jì)算）。

邊緣與云相反，意味著你是在本地運(yùn)行代碼(也就是說你能夠?qū)嶋H接觸到運(yùn)行代碼的設(shè)備)。在邊緣運(yùn)行代碼的主要優(yōu)點(diǎn)是沒有網(wǎng)絡(luò)延遲。由于物聯(lián)網(wǎng)設(shè)備通常要頻繁地生成數(shù)據(jù)，因此運(yùn)行在邊緣上的代碼非常適合基于物聯(lián)網(wǎng)的解決方案。

對比 CPU、GPU，深度剖析 TPU

TPU(Tensor Processing Unit, 張量處理器)是類似于CPU或GPU的一種處理器。不過，它們之間存在很大的差異。最大的區(qū)別是TPU是ASIC，即專用集成電路。ASIC經(jīng)過優(yōu)化，可以執(zhí)行特定類型的應(yīng)用程序。對于TPU來說，它的特定任務(wù)就是執(zhí)行神經(jīng)網(wǎng)絡(luò)中常用的乘積累加運(yùn)算。CPU和GPU并未針對特定類型的應(yīng)用程序進(jìn)行優(yōu)化，因此它們不是ASIC。

下面我們分別看看 CPU、GPU 和 TPU 如何使用各自的架構(gòu)執(zhí)行累積乘加運(yùn)算：

在 CPU 上進(jìn)行累積乘加運(yùn)算

CPU 通過從內(nèi)存中讀取每個(gè)輸入和權(quán)重，將它們與其 ALU (上圖中的計(jì)算器) 相乘，然后將它們寫回內(nèi)存中，最后將所有相乘的值相加，從而執(zhí)行乘積累加運(yùn)算。

現(xiàn)代 CPU 通過其每個(gè)內(nèi)核上的大量緩存、分支預(yù)測和高時(shí)鐘頻率得到增強(qiáng)。這些都有助于降低 CPU 的延遲。

GPU 上的乘積累加運(yùn)算

GPU 的原理類似，但它有成千上萬的 ALU 來執(zhí)行計(jì)算。計(jì)算可以在所有 ALU 上并行進(jìn)行。這被稱為 SIMD (單指令流多數(shù)據(jù)流)，一個(gè)很好的例子就是神經(jīng)網(wǎng)絡(luò)中的多重加法運(yùn)算。

然而，GPU 并不使用上述那些能夠降低延遲的功能。它還需要協(xié)調(diào)它的數(shù)千個(gè) ALU，這進(jìn)一步減少了延遲。

簡而言之，GPU 通過并行計(jì)算來大幅提高吞吐量，代價(jià)是延遲增加。或者換句話說:

CPU 是一個(gè)強(qiáng)大而訓(xùn)練有素的斯巴達(dá)戰(zhàn)士，而 GPU 就像一支龐大的農(nóng)民大軍，但農(nóng)民大軍可以打敗斯巴達(dá)戰(zhàn)士，因?yàn)樗麄內(nèi)硕唷?/p>

讀取 TPU 上的乘加操作的權(quán)重

TPU 的運(yùn)作方式非常不同。它的 ALU 是直接相互連接的，不需要使用內(nèi)存。它們可以直接提供傳遞信息，從而大大減少延遲。

從上圖中可以看出，神經(jīng)網(wǎng)絡(luò)的所有權(quán)重都被加載到 ALU 中。完成此操作后，神經(jīng)網(wǎng)絡(luò)的輸入將加載到這些 ALU 中以執(zhí)行乘積累加操作。這個(gè)過程如下圖所示：

TPU 上的乘加操作

如上圖所示，神經(jīng)網(wǎng)絡(luò)的所有輸入并不是同時(shí)插入 ALU 的，而是從左到右逐步地插入。這樣做是為了防止內(nèi)存訪問，因?yàn)?ALU 的輸出將傳播到下一個(gè) ALU。這都是通過脈動陣列 (systolic array) 的方式完成的，如下圖所示。

使用脈動陣列執(zhí)行乘加操作

上圖中的每個(gè)灰色單元表示 TPU 中的一個(gè) ALU (其中包含一個(gè)權(quán)重)。在 ALU 中，乘加操作是通過將 ALU 從頂部得到的輸入乘以它的權(quán)重，然后將它與從左編得到的值相加。此操作的結(jié)果將傳播到右側(cè)，繼續(xù)完成乘加操作。ALU 從頂部得到的輸入被傳播到底部，用于為神經(jīng)網(wǎng)絡(luò)層中的下一個(gè)神經(jīng)元執(zhí)行乘加操作。

在每一行的末尾，可以找到層中每個(gè)神經(jīng)元的乘加運(yùn)算的結(jié)果，而不需要在運(yùn)算之間使用內(nèi)存。

使用這種脈動陣列顯著提高了 Edge TPU 的性能。

Edge TPU 推理速度超過其他處理器架構(gòu)

TPU 還有一個(gè)重要步驟是量化 (quantization)。由于谷歌的 Edge TPU 使用 8 位權(quán)重進(jìn)行計(jì)算，而通常使用 32 位權(quán)重，所以我們應(yīng)該將權(quán)重從 32 位轉(zhuǎn)換為 8 位。這個(gè)過程叫做量化。

量化基本上是將更精確的 32 位數(shù)字近似到 8 位數(shù)字。這個(gè)過程如下圖所示：

量化

四舍五入會降低精度。然而，神經(jīng)網(wǎng)絡(luò)具有很好的泛化能力 (例如 dropout)，因此在使用量化時(shí)不會受到很大的影響，如下圖所示。

非量化模型與量化模型的精度

量化的優(yōu)勢更為顯著。它減少了計(jì)算量和內(nèi)存需求，從而提高了計(jì)算的能源效率。

Edge TPU 執(zhí)行推理的速度比任何其他處理器架構(gòu)都要快。它不僅速度更快，而且通過使用量化和更少的內(nèi)存操作，從而更加環(huán)保。