本教程詳細解釋了 Google TPU 脈動陣列加速卷積計算原理。
TPU 中計算卷積的方式和 GPU 不同，主要是依靠一種稱為 “脈動陣列” 的硬體電路結構來實作的。脈動陣列的主體部分是一個二維的滑動陣列，其中每一個節點都是一個脈動計算單元，可以在一個時鐘周期内完成一次乘加操作。脈動陣列中的行數和列數可以相等，也可以不相等，在每行每列的計算單元之間通過橫向或縱向的資料通路來實作資料的向右和向下滑動傳遞。
舉例：如下圖，特征圖像 F 大小為 3 x 3，通道數為 2；卷積核 W 和 G 大小為 2 x 2，卷積核個數為 2。
以這個例子來解釋脈動陣列實作卷積加速的過程。在這個例子中，采用固定卷積核權重，橫向和縱向脈動輸入特征值和中間部分和的方式。如下圖所示，首先将 2 個 2 通道 W 和 G 卷積核權重靜态地存儲到脈動陣列的計算單元中，其中對應同一個卷積核的 2 個通道權重被排列在同一列中。再将 2 個通道的輸入特征圖 F 排列展開，每一行都錯開一個時鐘周期，并準備依次輸入到脈動陣列中。
當要用 2 個通道輸入特征圖 F 和卷積核 W 及 G 進行卷積時，F 首先通過重排，通過左側數值載入器輸入到脈動陣列最左邊的一列中。輸入值從上到下，每一行進入脈動陣列的資料都會比上一行延遲一個時鐘周期。在第一個時鐘周期内，輸入值 F00 進入 W00 乘加單元與 W00 權重進行乘加運算，産生單元結果 Y00，此為 Y00 的第一部分和。脈動陣列第 1 個時鐘周期的計算狀态如下圖。
在第二個時鐘周期内，上一次 W00 乘加單元的部分和 Y00 被向下傳遞到 W01 乘加單元中，同時第二行輸入值 F01 與 W01 相乘的單元結果與傳遞下來的 Y00 相加得到 Y00第二次部分和。而 F01 進入 W00 乘加單元中，進行乘加運算求出 Y01 的第一次部分和。與此同時 F00 向右滑動進入 G00 乘加單元中求出單元結果 Y10 的第一次部分和。脈動陣列第 2 個時鐘周期的計算狀态如下圖。
以此類推，輸入的特征值沿着脈動陣列的行水準方向不斷向右滑動開啟不同卷積核各自的部分和。而對應于每一個卷積核的部分和沿着脈動陣列的列垂直方向不斷向下滑動并和目前計算單元的結果相累加，進而在每一列的最下方的計算單元中得到最終累加完成的對應于該卷積核所有通道的卷積結果。
這個例子中在第 8 個時鐘周期會得到第一個卷積結果，如下圖。
而從第 9 個時鐘周期開始，如下圖。每次都可以得到兩個卷積結果，直到所有的輸入特征圖像都被卷積完畢。脈動陣列通過橫、豎兩個方向的同時脈動傳遞，輸入值呈階梯狀進入陣列，順次産生每個卷積核的最終卷積結果。
脈動陣列的特性決定資料必須按照事先安排好的順序依次進入，是以每次要填滿整個矩陣就需要一定的啟動時間，而這段時間往往會造成硬體資源的浪費。啟動時間通常可以通過脈動陣列的 “行數 + 列數 - 1” 計算得到。當度過啟動時間後，整個脈動陣列進入滿載狀态，可以發揮出最大的吞吐率。
在上面的例子中采用了固定卷積核權重，橫向脈動輸入特征值，縱向脈動部分和的方式來計算卷積結果。同理，也可以采用固定部分和，橫向脈動輸入特征值，縱向脈動卷積核權重的方式來計算得到同樣的結果。事實上，可以任意選擇脈動三個變量中的兩個而固定另外一個的方式來實作卷積計算，具體的選擇會依據各種實作方式的實際情況和應用需求來決定。
需要指出的是，當計算的卷積核過大或者通道太多時，如果把所有通道的權重都串聯在一列上，那麼勢必造成脈動陣列過于龐大或者畸形，不便于現實的電路設計。這時候往往采用分割計算、末端累加的方式來解決這個問題。系統會把多通道的權重資料分割成幾個部分，每個部分都能夠适合脈動陣列的大小。然後依次對每個部分的權重進行計算，計算完成的結果會臨時存放在脈動陣列底部的一組加法累加器中。當換入另一組權重塊後結果就會被累加，直到所有權重塊都周遊完成，累加器累加所有的結果後輸出最終值。
TPU 中的脈動陣列僅僅完成了卷積的工作，而完成整個神經網絡的計算還需要其他計算單元的配合。如下圖，矢量計算單元通過輸入資料端口接收來自于脈動陣列中的卷積計算的結果，通過激活單元中的非線性函數電路來産生激活值，激活函數可以根據網絡需求來定制。在适量計算單元中，還可以通過歸一化函數對激活值進行歸一化，再經過池化單元就可以對激活值進行池化輸出。這些操作都由隊列子產品進行控制。例如隊列子產品可以通過配置參數的端口指定激活函數、歸一化函數或池化函數，以及處理步長等參數。矢量計算單元處理完成後将激活值發送到片上的統一緩沖區中暫存，作為下一層網絡的輸入。TPU 按照這種方法，一層一層地完成整個神經網絡的計算。
雖然脈動陣列主要用來加速神經網絡的推理計算，但這個架構不光是隻能處理卷積，它在進行通用矩陣計算時依然高效且強大，是以還可以用來完成除了卷積神經網絡以外的其他一系列工作，比如全連接配接神經網絡、線性回歸、邏輯回歸、分類（如K均值聚類）與視訊編碼和圖像處理等。
掃描下方二維碼即可關注我的微信公衆号【極智視界】，擷取更多AI經驗分享，讓我們用極緻+極客的心态來迎接AI ！