自动生成低精度深度学习算子
自動生成低精度深度學習算子
深度學習模型變得越來越大,越來越復雜,由于其有限的計算和能源預算,部署在低功耗電話和IoT設備上變得充滿挑戰。深度學習的最新趨勢是使用高度量化的模型,該模型可對輸入和幾位權重進行操作,諸如XNOR-Net,DoReFa-Net和HWGQ-Net之類的網絡正在穩步發展,提高了準確性。
下面是一個低精度圖形片段的示例。低精度卷積將量化的數據和位包放入適當的數據布局中,實現有效的比特卷積。輸出具有更高的精度,在對其進行重新量化,另一個低精度算子發送之前,將傳統的深度學習層(如批處理歸一化和ReLu)應用于該輸出。
低精度卷積管線。
理論上講,低精度算子比浮點算子使用更少的運算,可以實現巨大的加速。深度學習框架難以置信的良好優化的低級BLAS和LAPACK庫,利用了數十年的工程工作,CPU包含用于加速這些任務的內在指令。開發與8位量化甚至浮點算子的卷積之類的低級算子并不簡單。本文介紹了為CPU自動生成優化的低精度卷積的方法。聲明低精度算子,根據有效存儲的低精度輸入進行計算,描述用于描述實現參數搜索空間的調度。依靠AutoTVM來快速搜索空間,找到針對特定卷積,精度和后端的優化參數。
二進制計算背景
低精度模型的核心是biterial點積,使得僅使用按位運算和popcount即可計算卷積和密集算子。點積是兩個向量的元素逐次相乘,將所有元素相加來計算的,就像下面的簡單示例一樣。如果所有數據都是二進制數據,可以將輸入向量打包為單個整數,按位與打包的輸入,使用popcount對結果中的1進行計數來計算點積。根據輸入數據的量化方式,可以使用按位xnor代替按位二進制點積。
首先將輸入數據分成位平面,可以以這種方式計算任意精度的點積。一旦在該表示形式中,就可以對A和B的位平面之間的加權二進制點積求和來計算點積。二進制點積的數量隨A和B的精度乘積而增長,該方法僅適用于非常低精度的數據。
Bitserial點積
在TVM中定義算子
計算之前,需要對輸入數據進行位打包,以便可以訪問輸入數據的位平面,打包為受支持的數據類型,例如uint8或uint32。提供了一種靈活的位打包算子,該算子采用任意大小的輸入張量,返回一個位打包張量,用戶可以在其中指定位平面應位于哪個軸。
不同的位壓縮布局
一旦采用這種位打包格式,就可以按位計算低精度卷積。該數據沿輸入通道打包,位平面被添加到最里面的軸,數據打包為32位整數。比特卷積的計算與普通卷積類似,按位與(&)代替乘法,使用popcount累積打包數據中的值。位平面軸成為附加的歸約軸,計算輸入和內核的不同位平面之間的二進制點積。解壓縮格式和更高的精度計算輸出。
Input_bitpacked = bitpack(Input, activation_bits, pack_axis=3, bit_axis=4, pack_type=’uint32’)
Weights_bitpacked = bitpack(Filter, weight_bits, pack_axis=2, bit_axis=4, pack_type=’uint32’)
batch, in_height, in_width, in_channel_q, _ = Input_bitpacked.shape
kernel_h, kernel_w, _, num_filter, _ = Filter_bitpakced.shape
stride_h, stride_w = stride
pad_top, pad_left, pad_down, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
Computing the output shape
out_channel = num_filter
out_height = simplify((in_height - kernel_h + pad_top + pad_down) // stride_h + 1)
out_width = simplify((in_width - kernel_w + pad_left + pad_right) // stride_w + 1)
pad_before = [0, pad_top, pad_left, 0, 0]
pad_after = [0, pad_down, pad_right, 0, 0]
Input_padded = pad(Input_bitpacked, pad_before, pad_after, name=“PaddedInput”)
Treat the bitplane axes like additional reduction axes
rc = tvm.reduce_axis((0, in_channel_q), name=‘rc’)
ry = tvm.reduce_axis((0, kernel_h), name=‘ry’)
rx = tvm.reduce_axis((0, kernel_w), name=‘rx’)
ib = tvm.reduce_axis((0, input_bits), name=‘ib’)
wb = tvm.reduce_axis((0, weight_bits), name=‘wb’)
tvm.compute((batch, out_height, out_width, out_channel), lambda nn, yy, xx, ff:
tvm.sum(tvm.popcount(
Input_padded[nn, yy * stride_h + ry, xx * stride_w + rx, rc, ib] &
Weights_bitpacked[ry, rx, rc, ff, wb])) << (ib+wb))).astype(out_dtype),
axis=[rc, ry, rx, wb, ib]))
在調度中,應用了常見的優化,例如矢量化和內存切片,提供更好的內存局部性利用SIMD單元。其中的一些優化(如平鋪)需要針對特定的微體系結構調整參數。將這些參數作為旋鈕暴露給TVM,使用AutoTVM來自動同時自動調整所有參數。
可以制作小的微內核來替換最內層的計算循環,使用TVM的張量化原語調度。由于編譯器會產生次優的代碼,人們經常可以編寫較短的匯編序列,提高效率。這些微內核利用引入的新內在函數,幫助加速深度學習工作負載,使用巧妙的方式來改善內存訪問或減少所需的指令數量。
結果
樹莓派Raspberry
與16位整數TVM實現相比,Raspberry Pi 3B的卷積速度提高了。工作負載是ResNet18的卷積層。
與16位TVM實現相比,Raspberry Pi上低精度卷積的加速。
與移動設備上的高性能超低精度卷積的手動優化實現相比,Raspberry Pi 3B的2位激活,1位權重卷積加速得以實現,工作負載是ResNet18的卷積層。
手動優化實現2位權重1位激活Raspberry Pi卷積的速度。
x86
與32位浮點TVM實現相比,x86上的卷積速度有所提高。x86不支持此微體系結構的矢量化popcount,加速性較低。
與32位浮點TVM實現相比,x86低精度卷積的加速。
總結
以上是生活随笔為你收集整理的自动生成低精度深度学习算子的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 将TVM集成到PyTorch上
- 下一篇: TVM交叉编译和远程RPC