winograd卷积基本原理参考:
winograd把卷积拆分为4个步骤过程图示:
注意这张图里面隐藏了input和output channel。实际上每个空间维度里面还包含了batch和in/out channel维度。
winograd卷积优点: 通过特殊算法降低了计算量,从而获得加速效果。
winograd卷积缺点: 1,传统的卷积采用一个算子即可实现,而winograd分为4个算子,而input, weight transform后的数据量比原来卷积的输入大了很多倍,这使得这4个算子的访存量激增,但采用一个或多个算子融合这4个计算步骤的某几个非常困难。这个访存倍增的缺点使得winograd卷积不一定能够真正获得加速,只有在带宽相对较大而算力偏低的场景有效,而带宽不足,算力很大的硬件上,winograd不见得能加速。2,winograd卷积数值精度比传统卷积和img2col算法差,tile size越大,理论加速比越高,但数值精度越差,可能导致结果不可用。
输入数据格式为[n, h, w, c],input transform后格式原始的论文[Fast Algorithms for Convolutional Neural Networks]的格式是[ho/2*w0/2, 16, n, ci],实际也常用[n, ho/2*w0/2, 16, 1, ci]。ho, wo是输出height和width尺寸,除以2是因为两个相邻元素共享一个4x4矩阵。
输入变换是空间维乘以变换矩阵,由于每个空间维元素实际上对应于一个2维向量,实际上相当于每个空间维元素对应的的向量之间做各种乘加运算。如下图。
weight原始格式为[c0, ci, h, w],这里h = w = 3, 先处理为[h, w, ci, co], 然后weight transform转换后为[4*4, ci, co]。
转换后的输入和weight做batch矩阵乘得到数据格式为[n, ho/2*w0/2, 4*4, 1, co]。
最后做输出transform得到格式为[n, ho/2*w0/2, 2*2, co],也就是[n, ho, wo, co]
输入变换计算过程,假设输入x的4x4空间维内容标记为a-p(可以借助Matlab符号计算来得到转换结果矩阵与x的计算关系):
transform输入到输出公式
filter_trans =
[ a, a/2 + b/2 + c/2, a/2 - b/2 + c/2, c]
[ a/2 + d/2 + g/2, a/4 + b/4 + c/4 + d/4 + e/4 + f/4 + g/4 + h/4 + i/4, a/4 - b/4 + c/4 + d/4 - e/4 + f/4 + g/4 - h/4 + i/4, c/2 + f/2 + i/2]
[ a/2 - d/2 + g/2, a/4 + b/4 + c/4 - d/4 - e/4 - f/4 + g/4 + h/4 + i/4, a/4 - b/4 + c/4 - d/4 + e/4 - f/4 + g/4 - h/4 + i/4, c/2 - f/2 + i/2]
[ g, g/2 + h/2 + i/2, g/2 - h/2 + i/2, i]
data_trans =
[ a - c - i + k, b + c - j - k, c - b + j - k, b - d - j + l]
[ e - g + i - k, f + g + j + k, g - f - j + k, f - h + j - l]
[ g - e + i - k, j - g - f + k, f - g - j + k, h - f + j - l]
[ e - g - m + o, f + g - n - o, g - f + n - o, f - h - n + p]
out_trans =
[ a + b + c + e + f + g + i + j + k, b - c - d + f - g - h + j - k - l]
[ e + f + g - i - j - k - m - n - o, f - g - h - j + k + l - n + o + p]
可以用另一种视角来看待这个转换过程:
# result = (x.reshape([1, -1]) * data_trans).reshape([4, 4]),wegit, out trans以此类推
data_trans =
[
a* 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
b* 0 1 -1 1 0 0 0 0 0 0 0 0 0 0 0 0
c* -1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0
d* 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0
e* 0 0 0 0 1 0 0 0 -1 0 0 0 1 0 0 0
f* 0 0 0 0 0 1 -1 1 0 -1 1 -1 0 1 -1 1
g* 0 0 0 0 -1 1 1 0 1 -1 -1 0 -1 1 1 0
h* 0 0 0 0 0 0 0 -1 0 0 0 1 0 0 0 -1
i* -1 0 0 0 1 0 0 0 1 0 0 0 0 0 0 0
j* 0 -1 1 -1 0 1 -1 1 0 1 -1 1 0 0 0 0
k* 1 -1 -1 0 -1 1 1 0 -1 1 1 0 0 0 0 0
l* 0 0 0 1 0 0 0 -1 0 0 0 -1 0 0 0 0
m* 0 0 0 0 0 0 0 0 0 0 0 0 -1 0 0 0
n* 0 0 0 0 0 0 0 0 0 0 0 0 0 -1 1 -1
o* 0 0 0 0 0 0 0 0 0 0 0 0 1 -1 -1 0
p* 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1
]
weight_trans =
[
a* 4 2 2 0 2 1 1 0 2 1 1 0 0 0 0 0
b* 0 2 -2 0 0 1 -1 0 0 1 -1 0 0 0 0 0
c* 0 2 2 4 0 1 1 2 0 1 1 2 0 0 0 0
d* 0 0 0 0 2 1 1 0 -2 -1 -1 0 0 0 0 0
e* 0 0 0 0 0 1 -1 0 0 -1 1 0 0 0 0 0
f* 0 0 0 0 0 1 1 2 0 -1 -1 -2 0 0 0 0
g* 0 0 0 0 2 1 1 0 2 1 1 0 4 2 2 0
h* 0 0 0 0 0 1 -1 0 0 1 -1 0 0 2 -2 0
i* 0 0 0 0 0 1 1 2 0 1 1 2 0 2 2 4
]/4
out_trans =
[
a* 1 0 0 0
b* 1 1 0 0
c* 1 -1 0 0
d* 0 -1 0 0
e* 1 0 1 0
f* 1 1 1 1
g* 1 -1 1 -1
h* 0 -1 0 -1
i* 1 0 -1 0
j* 1 1 -1 -1
k* 1 -1 -1 1
l* 0 -1 0 1
m* 0 0 -1 0
n* 0 0 -1 -1
o* 0 0 -1 1
p* 0 0 0 1
]
可以看到这里面非常有规律,可以利用这个规律来显著加速winograd计算过程。
端侧推理引擎如mnn由于通常没有或只有很少的shared mem,线程之间数据交换能力比较弱,通常按照这个流程把整个计算流程拆分为了三个步骤而不是写在一个完整的kernel:input transform, matmul, output transform。而权重在推理时是常量,可以通过常量折叠提前做好weight transform。
对于NHWC的数据,通常采用每个thread读取4x4,并且channel深度为4的输入数据做input transform。
矩阵乘部分shape是[n, ho/2*w0/2, 16, 1, ci] * [16, ci, co],最内层的[1, ci] * [ci, co]可以考虑通过每个线程计算1x1*1x4 tile大小矩阵乘(与常用的4x1*1x4或8x1*1x8 tile原理相同,参考[施工中] CUDA GEMM 理论性能分析与 kernel 优化 - 知乎)。但这个矩阵乘shape太小,导致每个线程计算量小,而且需要创建太多线程。
另一种做法是input transform输出的shape格式是[n, 16, ho/2*w0/2, ci],然后与weight transform后的[16, ci, co]做矩阵乘,这样最内层的矩阵乘大小显著增大为[ho/2*w0/2, ci]*[ci, co],更有利于性能优化。矩阵乘输出格式为[n, 16, ho/2*w0/2, co],output transform变成[n, ho/2*w0/2, 2*2, co]。相当于在input和output transform里面包含一个transpose的操作。
跟卷积通过im2row+matmul的实现相比,由于相邻的卷积框有重复的数据,对于3x3 stride=1,im2row每个位置数据读取一次,写出9次,因此im2row写回和matmul读取数据量增大了9倍。
winograd input transform相当于4x4 kernel, stride = 2,input transform写回和matmul读取数据数据量相比输入增大4倍。
思考:相比NHWC或者NCHW输入格式,NCHW4输入格式是否对winograd卷积会带来什么帮助?
本文中的im2row是相对于im2col而言,im2col是把一个卷积核覆盖的输入数据展开成矩阵的一列作为矩阵乘的输入1,而权重展开是矩阵乘的输入0。而im2row是把卷积核覆盖的输入数据展开成矩阵的一行作为矩阵乘的输入0,而权重展开是矩阵乘的输入1。im2row方式连续地址写回性能更好,并且推理引擎通常认为矩阵乘的input 1而不是input 0是常量。当然im2col得到的输入数据展开后作为矩阵乘的输入1也不是绝对的,也可以直接作为矩阵乘输入0并标记trans_a=1, 因为矩阵乘的优化通常也把输入0转置一下更有利于性能优化(不同于转置b更好的老旧的"常识")。
其他文章参考:
MegEngine Inference 卷积优化之 Im2col 和 winograd 优化_旷视的博客-优快云博客