Resize算子替换加速
背景
源自工作中出现的一个问题:模型在导出编译之后,有两个2x Resize算子cpu fallback了,一个是Bilinear Resize,一个是Nearest Resize。解决方法是使用Conv和Matmul算子进行替代。Conv和Matmul走的是BPU计算的路径,完美解决了cpu fallback的问题。
本文从两种不同Resize的计算方式入手,详细解释为什么能用Conv和Matmul算子替代。需要注意的是,仅针对2x Resize,其他scale不一定满足这一性质。
Resize如何计算?
Reisze的作用,就是对一张图像进行缩放。
假设输入:
2 x 2 image
10 20
30 40
放大 2 倍,变成4x4:
? ? ? ?
? ? ? ?
? ? ? ?
? ? ? ?
其中原来的像素会保留:
10 20
30 40
其他的像素网格该怎么办?需要插值。
放大后的图像与原图如何对齐呢?这里就要介绍align_corners了。
align_corners
align_corners可以简单理解为:
放大以后,原图的角点像素,要不要和新图的角点严格对齐?

可以看到,在align_corners=True的情况下,四个角点都是对其的,而False的时候没有对齐。
Bilinear Resize 2x
在使用Bilinear方式插值的情况下:当align_corners=False时,采用conv替换;align_corner=True时,采用Matmul替换。下面来详细解释原因:
align_corners=Flase:Conv替换
先说如何替换成Conv计算:
源图(N,C,H,W) -> Conv2d(kernel=3x3, padding=1, 4C 输出通道)
-> (N,4C,H,W) -> PixelShuffle(20) -> 2x放大图(N,C,2H,2W)
输出图的第i个像素,对应到源图的位置(源坐标)是:
y_src(i) = (i + 0.5) x step - 0.5
其中step=输入尺寸/输出尺寸。当放大两倍时,step=0.5
output:
i:
0 1 2 3 4 ...
对应src:
-0.25 0.25 0.75 1.25 ...
输出像素不断落在输入像素之间。

蓝色的是缩放后的坐标,黑色的是源坐标。
但是,这个映射公式得到的源坐标是一个小数,源图上不存在位置坐标为小数的像素点。于是,就需要用Bilinear来插值了。

红色的坐标代表缩放后的图在源图取像素的位置,黑色的坐标代表源图上的像素点(整数位置)
- 纵向:y_src=1.25,离y=1距离0.25,所以给y=1那一行的权重为1-0.25=0.75(距离越小,权重越大);离y=2距离0.75,所以给y=2那一行的权重为1-0.75=0.25。
- 横向:x_src=2.75,离x=2距离0.75,所以给x=2那一列的权重为1-0.75=0.25;离x=3距离0.25,所以给x=3那一列的权重为1-9.25-0.75。
所以我们可以得到:
| 角 | 权重 |
|---|---|
| A | 0.75x0.25=0.1875 |
| B | 0.75x0.75=0.5625 |
| C | 0.25x0.25=0.0625 |
| D | 0.25x0.75=0.1875 |
可以得到最终的output:
output = 0.1875xA + 0.5625xB + 0.0625xC + 0.1875xD
用固定的四个权重去乘变化的四个角点,这不就是卷积运算吗?这四个权重不就是卷积核的weight吗?

刚刚我们讨论的是第一象限。映射源坐标的位置不外乎这四个点,也就是这四个象限。每个象限的点对应ABCD四个角点的权重都不同,那我们是不是只要用四个卷积核去表示就好了?于是,这个卷积核的输出通道维度就是4,我们只需要把这四个象限对应的权重写死到卷积核的四个通道里就行了。
那卷积核的大小如何确定呢?根据公式:y_src(i) = (i + 0.5) x step - 0.5,当i=0时,y_src=-0.25,也就是超出了源图的左边界。所以我们要把padding设置成1,kernel size=3x3。
卷积运算之后,得到的图片维度是(N,4C,H,W),我们还需要进行PixelShuffle(2)操作,把4C均摊到H和W维度上,于是就变成了(N,C,2H,2W),也就是放大两倍后的图像。完美!
align_corners=True:Matmul替换
这时,源坐标的计算公式就与False时不同了:
y_src(i) = i x (H_in - 1) / (H_out - 1)
False时,step=H_in / H_out。不管是2x2 -> 4x4,还是3x3 -> 6x6,step恒定为0.5
True时就不一样了,step的大小是随着输入图像的尺寸变化的,这就导致卷积核是不固定的,也就不能使用Conv来替换了。如何用Matmul来替换呢?我们继续推导:
我们考虑1D的情况。y_src(i)是一个小数,总会落在两个整数之间。我们设:
lower(i) = floor(y_src(i))
upper(i) = lower(i) + 1
frac(i) = y_src(i) - lower(i)
按照Bilinear的规则去加权:
input[lower(i)]的权重 = 1 - frac(i)
input[upper(i)]的权重 = frac(i)
所以,最终输出就是:
output = input[lower(i)] * (1 - frac(i)) + input[upper(i)] * frac(i)
上面这个公式其实就是一个稀疏矩阵乘法。我们构造这样一个矩阵M:

output = input @ M
每个输出位置 = input沿这一维的加权求和,权重来自M的列。
以上推导的是一维的情况,二维只需要做两次Matmul即可:先沿W方向,再沿H方向
总体流程:
step1:沿W方向做matmul:
input @ M_w
(N,C,H_in,W_in) @ (W_in,W_out) -> (N,C,H_in,W_out)
step2:先transpose,再沿H维做matmul,再transpose回来
(N,C,H_in,W_out) -> (N,C,W_out,H_in) @ (H_in,H_out) -> (N,C,W_out,H_out) -> (N,C,H_out,W_out)
于是,Resize就变成了两次Matmul+两次transpose
Nearest Resize 2x:Conv替换
Nearest Resize就简单得多。它就是找最近的一个像素,直接复制。
例如:
input:
1 2
3 4
output:
1 1 2 2
1 1 2 2
3 3 4 4
3 3 4 4
对于2x Nearest,align_corners的取值不影响计算结果。
使用Conv替换的计算流程与Bilinear Resize的完全相同:
源图(N,C,H,W) -> Conv2d(kernel=3x3, padding=1, 4C 输出通道)
-> (N,4C,H,W) -> PixelShuffle(20) -> 2x放大图(N,C,2H,2W)
区别在于卷积核的选取。3x3的Conv kernel变化:
0 0 0
0 1 0
0 0 0
只是单纯的复制。