歡迎來到AI的張量世界。是時候直面維度災難了

全文共

4429

字，預計學習時長

9

分鐘

圖片來源：pexels。com/@pixabay

16並不是一個很大的數字。那麼，在64個時鐘週期內，計算一個卷積神經網路（Convolutional Neural Network，簡稱CNN），使其能夠從16個輸入通道、3*3張量卷積中生成具有16*16張瓦片圖（tile）的16個輸出通道，將需要多少MAC（Multiply-Accumulate Unit，乘加器）呢？

如果不使用快速演算法，那麼答案將會是至少9216個。通常，9216個MAC可以用來構建一個96*96的脈衝陣列，但是計算一個96*96的矩陣乘法（matrix multiplication，簡稱MM）將延遲至少96個時鐘週期。也就是說，構建這個脈動陣列需要96個96*96矩陣乘法。

歡迎來到AI的張量世界。是時候直面維度災難了。

張量vs矩陣

本文標題靈感來自Charles F。Van Loan教授在2010年所作的關於張量的模展開矩陣的演講中的一句話——所有張量暗地裡都希望它們是矩陣！

這一陳述表明了張量研究領域想要先把張量展開成矩陣，然後利用成熟的矩陣理論框架對其進行研究。同時，儘管主流AI框架認為張量是最基本的資料型別，但在行業標準實踐中，也常將張量展開為矩陣，以便利用高度最佳化的矩陣乘法（MM）庫和矩陣乘法加速器（MMAs）庫。通常，AI領域認為矩陣是特殊的張量。

一些主流的看法認為：

1。 可以利用MM中的並行性和資料共享模式。

2。 和張量相比，矩陣更適用於可程式設計硬體。

3。 MM有原生硬體加速方法，即脈動陣列。

然而，

1。 從結構上來說，CNN和MM是相同的。展開張量來利用和MM相同的並行性和資料共享模式是完全沒有必要的。

2。 在從儲存層次的下層向上層執行的過程中，由於時間區域性性，矩陣會遞迴式分塊；由於空間區域性性，矩陣會壓縮打包。最終，矩陣會變成微面板，即小塊行或列，併為軟體微核心或GPU著色核心所用。

3。 在谷歌TPU v1：256*256脈衝陣列概述中，筆者解決了方形脈動陣列的可擴充套件性問題。自此，使用多個小型脈動陣列似乎成了主流。因此，在矩陣成為能夠為脈動陣列所用的方形矩陣前，仍然需要被分塊和壓縮打包。

所以，由展開張量得到的矩陣實際上被分塊和壓縮打包成合適的結構，用於高效能執行（如下圖所示）。前者被稱為平滑展開，後者為分塊展開。由於數十年的研究、構建和利用分塊矩陣框架用於高效能MM執行，矩陣成為CNN和AI中最常見的預設資料型別。

MM的分塊展開

主流看法認為，CNN的特徵圖可看作是某個矩陣中的一列，卷積核（filter）則可展開成為一列中的某些連續矩陣元素。由於平滑展開，特徵圖中的相鄰畫素值被空間和時間重複引用的功能被抹去了。

拯救分塊張量

MM中的遞迴矩陣分塊展開有一個重要特徵，即當矩陣逐漸接近硬體的裸金屬（bare metal）時，高水平的時空位置得以保留。那麼，考慮CNN中的資料位置會不會像張量形式中的位置一樣被保留將會非常有趣？

將特徵圖分成瓦片圖，並保持其平鋪結構，以便於重複利用卷積核、將快速演算法用於瓦片圖以及達到細粒度的SIMD（Single Instruction Multiple Data，單指令多資料流）並行是一個合理的做法。當接近裸金屬時，張量應該保持張量的形式，並保持平鋪結構和完整特徵圖中的資料區域性性。

此外，必須解決輸入特徵圖和輸出特徵圖中都存在的區域性模式。前者使得多個輸入特徵圖的資料能夠在計算多個輸出特徵圖時實現共享。後者則使得更多方面能夠共享輸入特徵圖。不同之處在於，由於生成輸出畫素值並不需要用到所有的輸入特徵圖，所以不能共享所有的特徵圖。更何況，把所有特徵圖放在晶片上並不實際。總的來說，所有維度都要按一定程度進行分塊，這使得張量能夠被分成更小的單元，即分塊張量。

分塊張量即介面為張量的張量，和分塊矩陣相對應。和分塊矩陣被高效能計算（high-performance computing，簡稱HPC）機群用於MM相似，分塊張量可用於解決維度災難和保留CNN的資料區域性性。張量包，相當於微通道或MM中的方形子矩陣，是最基本的張量單元。它必須按照原子級執行，以利用所有維度的空間區域性性。由張量包構成的張量塊也是一種張量單元。它必須在整體計算單位和外部記憶之間轉移，以促進張量包之間的時間區域性性。

原子級張量包執行可根據最小充分輸入通道量來生成具有最小充分大小瓦片圖的最小充分輸出通道量。由於張量中的維度災難，即使在每個維度的張量包都很小時，上述張量包執行也能發揮很大作用。它可以在張量塊中迭代或並行執行，來解決更嚴峻的問題。該方法將在下文中半正式地詳細闡述。

CNN中的分塊張量

輸入為A、輸出為C的CNN包括多個輸入特徵圖（input feature maps，簡稱IFMs）和多個輸出特徵圖（output feature maps，簡稱OFMs）。這些可以看作是三維張量，用x，y表示特徵圖，用輸入深度w索引IFMs，用輸出深度z索引OFMs。為了達到細粒度SIMD並行和利用具有特殊區域性性的快速演算法，每個特徵圖將會進一步沿著x維和y維分成瓦片圖。相應的索引（tx， ty）分別表示輸入瓦片圖和輸出瓦片圖。對於每一對唯一的IFM w和OFM z，有卷積核B（w， z），通常由3*3的卷積核引數構成。輸入張量和輸出張量成為瓦片圖的分塊張量，如下圖所示：

CNN中的瓦片圖分塊張量

利用張量理論符號，可以更加完整、準確地表示平鋪CNN：

冒號表示採用該維度中的所有資料。如A（：， ：， w）表示採用了IFM w中的所有瓦片圖，可表示IFM w。

CNNs在結構上和MMs相同

主流想法認為必須把張量平滑展開為矩陣，以便於利用MMs的並行性和資料共享模式。然而，這實際上會適得其反。如下圖所示，從結構上來說，CNNs和MMs的並行性和資料共享模式是相同的。這也是MMs在CNNs中廣泛使用的原因。

CNNs在結構上和MMs相同

因為C行是獨立計算的，所以A行可以被劃分。同理，因為畫素值可以被單獨計算，所以特徵圖可以被分成瓦片圖。

瓦片圖中和MM相同的並行性和資料共享模式保持完整。

在同一特徵圖中，輸出瓦片圖的卷積核相同，但輸入瓦片圖不同。在輸出維度中，輸出瓦片圖的輸入瓦片圖相同，但卷積核不同。下表利用張量理論符號更加完整地解釋了MMs和CNNs的相同之處。

深度也需要被劃分

現在，假設晶片的記憶容量是有限的，並且需要基於晶片上的一大堆資料，在IFMs、OFMs以及卷積核中充分發揮CNNs的與MM相同的資料共享模式。這些資料的幾何結構會是怎樣的呢？如下圖所示，這需要有充分的（x， y）來重複利用卷積核，充分的w使得有足夠的輸入資料被共享，以及充分的z來有效共享輸入資料：

晶片緩衝器必須有充分的資訊範圍涵蓋所

從上圖中可以觀察到，需要拆分w維，z維，x維以及y維。這一步已經完成了。這一步是為了保證大量的輸入資料和輸出資料會分別涵蓋最小充分IFMs量和最小充分OFMs量。引入新的索引tw和tz，分別用於定義IFMs組和OFMs組。在IFM組tw中，定義一個輸入張量包A（tx， ty， tw），作為一組對齊的瓦片圖（tx， ty）。同理，在OFM組tz中，定義C（tx， ty， tz）作為一組對齊的瓦片圖（tx， ty）。B（tw， tz）表示一組卷積核，每個卷積核用於由一個來自組tw的IFM和一個來自組tz的OFM構成的一對。這就是卷積核張量包。輸入，輸出和卷積核張量構成張量包中的分塊張量，如下圖所示：

CNN中張量包中的分塊張量

張量包中和MM相同的並行性和資料共享模式保持完整。

軟體中的張量微核或者硬體中的張量包引擎可以用於處理卷積一個輸入張量包A（tx， ty， tw）和一個卷積核張量包B（tw， tz）的原子級操作，如下圖所示：

原子級張量包操作

假設輸入和輸出瓦片圖分別為6*6和4*4，並把IFM組和OFM組大小設定為8。那麼在不使用快速演算法的情況下，執行一次3*3卷積，需要2304個MAC和4個時鐘週期。這2304條並行執行平均分佈在所有維度中，包括特徵圖中的x維和y維、輸入深度w和輸出深度z。而如果將特徵圖分為瓦片圖，如4*4規模，那麼快速演算法，如Winograd就可以得到應用，則這2304條並行執行僅需576個MAC就可以完成了。

張量遞迴劃分

在計算單元中，張量包是基本單元。為了保留tile函式間的資料位置和平鋪結構，在完整張量、張量包和張量塊之間引入一箇中間塊，以完整涵蓋想要放到晶片上的張量包。張量包是tile函式間的時間位置的基本單元，也是當有充分計算單元、晶片記憶容量和供應頻寬時的空間區域性性的基本單元。在x，y和z維上，一個張量塊可定義為（bx， by， bz）。

下文展示了一個張量遞迴劃分的例子。整個張量是一個分塊張量，該分塊張量可分成4*4*2個分塊張量，每一塊可再進一步分成1*1*8個分塊張量，每一塊包含4*4個瓦片圖。其中（dx， dy）是分支，用於找到每個瓦片圖中的畫素值，dz是索引，用於定義每個組中的特殊特徵圖。

張量遞迴劃分，用於保留特徵圖中的區域性

張量塊在儲存器中線性儲存。但其儲存順序並不需要在意。對於一個張量塊，有兩種不同的分塊展開順序，一種不斷最佳化將張量塊儲存到DRAM（Dynamic Random Access Memory，動態隨機存取儲存器）中，另一種則不斷最佳化將張量塊展示到計算單元中。軟體和/或硬體中的換位機制需要在執行過程中從一種格式轉換到另一種格式，如下圖所示。

張量塊換位

由換位導致的延遲可以由雙緩衝來解決。當一個張量塊被展示到計算單元中時，它被展開成張量包，並按（tx， ty）排序，如此一來，張量包就可以按相應的計算單元在（x， y）域並行執行。

所有張量暗地裡都想做自己

Van Loan教授在一場關於分塊張量的分塊展開的談話中講到：

分塊展開保留了資料結構和區域性性……我認為，分塊在張量計算中的影響力最終將會和其在矩陣計算中的影響力一樣大。

在多維時間區域性性中，張量被分塊展開成張量塊，以便於晶片上不同維度的資料實現共享。張量塊則可以進一步分塊展開成張量包，每個張量包有各個維度的最小充分資訊範圍。在本文所舉的例子中，執行這樣的張量包將會需要2304條並行執行。在張量塊中，和MM相同的並行性和資料共享模式應用為張量包。

就本文所瞭解的來看，本文可能在計算機史上第一次討論了將分塊張量的概念應用到裸金屬中。跨x維和y維的更好的、規模化的並行性可用於處理高畫質影片/影象。基於CNN革命性的能力和其在張量中的深刻根基，CNN可能成為第一個分塊張量殺手鐧。

所有張量暗地裡都想要成為自己，即分塊張量，並釋放潛能，在AI硬體領域實現突破。

留言 點贊 關注

我們一起分享AI學習與發展的乾貨

歡迎關注全平臺AI垂類自媒體 “讀芯術”