深度学习算子优化-FFT

背景

在数字信号和数字图像领域，对频域的研究是一个重要分支。

我们日常“加工”的图像都是像素级，被称为是图像的空域数据。空域数据表征我们“可读”的细节。如果我们将同一张图像视为信号，进行频谱分析，可以得到图像的频域数据。观察下面这组图，频域图中的亮点为低频信号，代表图像的大部分能量，也就是图像的主体信息。暗点为高频信号，代表图像的边缘和噪声。从组图可以看出，Degraded Goofy 与 Goofy 相比，近似的低频信号保留住了 Goofy 的“轮廓”，而其高频信号的增加使得背景噪点更加明显。频域分析使我们可以了解图像的组成，进而做更多的抽象分析和细节处理。

Goofy and Degraded Goofy

实现图像空域和频域转换的工具，就是傅立叶变换。由于图像数据在空间上是离散的，我们使用傅立叶变换的离散形式 DFT（Discrete Fourier Transform）及其逆变换 IDFT（Inverse Discrete Fourier Transform)。Cooley-Tuckey在DFT的基础上，开发了更快的算法 FFT（Fast Fourier Transform）。

DFT/FFT在数字图像领域还有一些延伸应用。比如基于 DFT 的 DCT（Discrete Cosine Transform，离散余弦变换）就用在了图像压缩JPEG算法和图像水印算法。

JPEG 编码是通过色彩空间转换、抽样分块、DCT 变换、量化编码实现的。其中 DCT 变换的使用将图像低频信息和高频信息区分开，在量化编码过程中压缩了少量低频信息、大量高频信息从而获得尺寸上压缩。从猫脸图上可看出随着压缩比增大画质会变差，但是主体信息还是得以保留。

猫脸不同 jpeg 画质（压缩比）

图像水印算法通过 DCT 将原图转换至频域，选取合适的位置嵌入水印图像信息，并通过 IDCT 转换回原图。这样对原图像的改变较小不易察觉，且水印通过操作可以被提取。

DFT/FFT 在深度学习领域也有延伸应用。比如利用 FFT 可以降低卷积计算量的特点，FFT\_Conv 算法也成为常见的深度学习卷积算法。本文我们就来探究一下频域算法的原理和优化策略。

DFT的原理及优化

公式

无论是多维的 DFT 运算，还是有基于 DFT 的 DCT/FFT\_Conv, 底层的计算单元都是 DFT\_1D。因此，DFT\_1D 的优化是整个FFT类算子优化的基础。

DFT\_1D 的计算公式：

单位复根的性质

Cooley-Tuckey FFT算法

N=8 碟形算法图

N=8 的计算序列被分成了 3 级，每一级 (stage) 有一个或多个块 (section)，每个块中包含了一个或者多个蝶形（butterfly) , 蝶形的计算就是 DFT 运算的 kernel。

每一个 stage 的计算顺序：

取输入
乘以转换因子
for section_num, for butterfly_num，执行radixN_kernel
写入输出

看 N=8 的蝶形算法图，stage = 1 时，运算被分成了 4 个 section，每个section的 butterfly_num = 1 。stage = 2 时， section_num = 2，butterfly_num = 2。stage = 3时， section_num = 1， butterfly_num = 4。

可以观察到，从左到右过程中 section_num 不断减少， butterfly_num 不断增加，蝶形群在“变大变密”，然而每一级总的碟形次数是不变的。

实际上，对于长度为 N ，radix = r 的算法，我们可以推得到：

radix-2 的幂次优化

radix-非2的幂次优化

其他优化

上述两个章节描述的是 DFT\_1D 的通用优化，在此基础上还可以做更细致的优化，可以参考本文引用的论文。

对于全实数输入的，由于输入的虚部为 0，进行旋转因子以及radix\_N\_kernel 的复数运算时会有多余的运算和多余的存储，可以利用 split r2c 算法，视为长度为 N/2 的复数序列，计算 DFT 结果并进行 split操作得到 N 长实数序列的结果。
对于 radix-2 的幂次算法，重新计算每个 stage 的输入/输出 stride 以取消第一级的位元翻转可以进一步减少访存的开销。
对于 radix-N 算法，在混合基框架下 N = radix1^m1 * radix2^m2，合并较小的 radix 为大的 radix 以减少 stage。

DFT 延展算法的原理及优化

DCT 和FFT\_conv 两个典型的基于 DFT 延展的算法，DFT\_1D/2D 的优化可以很好的用在这类算法中。

DCT

FFT\_conv

Conv 是深度学习最常见的运算，计算conv常用的方法有 IMG2COL+GEMM， Winograd， FFT\_conv。三种算法都有各自的使用场景。

FFT_conv 算法的流程：

将 Feature Map 和 Kernel 都 zero-pad 到同一个尺寸,进行 DFT 转换。
矩阵点乘
将计算结果通过 IDFT 计算出结果。

该算法将卷积转换成点乘，算法复杂度是 O(nlogn),小于卷积的 O(n^2), 在输入的尺寸比较大时可以减少运算量，适用于大 kernel 的 conv 算法。

深度学习计算中， Kernel 的尺寸要远小于 Feature Map，因此 FFT\_conv第一步的 zero-padding 会有很大的开销，参考论文2里提到可以通过对 Feature map进行分块，分块后的 Feature Map 和 Kernel 需要 padding 到的尺寸较小，可以大幅减小这一部分的开销。优化后 fft\_conv 的计算流程为：