Python Keras/Tensorflow——conv2d的傅里叶逐点乘法实现，运行速度比空间卷积慢4倍

根据卷积定理，卷积变为傅里叶域中的逐点乘法，在许多以前的工作中，由于将卷积运算转换为逐点乘法运算，傅里叶变换的开销被增益所掩盖，例如：

为了复制这一点，我尝试将keras.layers.Conv2D（）层替换为一个接受输入数据rfft的自定义层（我在将数据输入到模型之前获取rfft以减少训练时间），初始化“no_of_kernels”与图像大小相同的核数，获取其rfft，将输入和内核逐点相乘并返回乘积（是的，不采用irfft，因为我想在傅立叶域中进一步训练网络本身）——

在该层中，调用函数的实现如下- 注意-在我的数据集中，即MNIST image height=width，因此转置工作正常

def call(self, x):
        fft_x = x #(batch_size, height, width, in_channels)
        fft_kernel = tf.spectral.rfft2d(self.kernel) #(in_channels, height, width, out_channels)
        fft_kernel = tf.transpose(fft_kernel, perm=[2, 1, 0, 3]) #(width, height, in_channels, out_channels)
        output  = tf.einsum('ijkl,jklo->ijko', fft_x, fft_kernel)
        return output 

该代码保留了Keras Conv2D layer给出的精度，但其运行速度比Conv2D慢4倍左右，因此无法实现转换为傅里叶域的目的。谁能解释一下为什么会发生这种情况，以及我如何在傅里叶域中复制快速卷积的结果

（注意——对于那些可能觉得tf.spectral.rfft2d（self.kernel）可能是开销的人来说，我已经证实了这一点

此外，我认为Conv2D函数可能会将4D输入张量和核展平，以将其简化为此处解释的矩阵乘法-。我想不出任何智能的展平方法等来执行逐点乘法，除了将其视为点积，就像我使用tf.einsum所做的那样。是否有智能方法做点式乘法吗？） 谢谢

编辑- 整个层的实现供参考-

class Fourier_Conv2D(Layer):
    def __init__(self, no_of_kernels, **kwargs):
        self.no_of_kernels = no_of_kernels
        super(Fourier_Conv2D, self).__init__(**kwargs)

    def build(self, input_shape):
        self.kernel_shape = (int(input_shape[3]), int(input_shape[1]), int(input_shape[2]), self.no_of_kernels)
        self.kernel = self.add_weight(name = 'kernel', 
                                      shape = self.kernel_shape, 
                                      initializer = 'uniform', trainable = True)
        super(Fourier_Conv2D, self).build(input_shape)

    def call(self, x):
        fft_x = x
        fft_kernel = tf.spectral.rfft2d(self.kernel)
        fft_kernel = tf.transpose(fft_kernel, perm=[2, 1, 0, 3])
        output  = tf.einsum('ijkl,jklo->ijko', fft_x, fft_kernel)
        return output       

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1], input_shape[2], int(self.no_of_kernels/2)+1)

---

我认为您的结果一点也不令人惊讶，在Keras中Conv2D的实现留给后端，大多数后端（如TensorFlow）都有非常优化的卷积运算版本，特别是如果您使用CuDNN。因此，您自己的版本应该比简单的实现快，但比高度优化的版本慢

---

为了进行有意义的比较，您可能需要实现一个基线Conv2D，它以一种简单的方式进行卷积，而不进行任何优化。

是的，谢谢。我将尝试实现一个基本的二维空间实现，然后进行比较。但我没有使用CuDNN，所以这两个代码都在CPU上运行。Tensorflow后端是否也在CPU上进行了优化？是的，它也应该进行优化，可能使用im2col和BLAS。好的，谢谢。我读到Tensorflow使用Eigen优化CPU操作。因此，我是否需要使用Eigen等编写一个自定义操作，以使fourier_conv2d与conv2d相比/比conv2d更快？在conv2d层中，我使用了一个3x3内核。在fourier_conv2D中，我初始化了一个与图像大小相同的内核。由于这增加了可学习参数的数量，我还尝试初始化3x3内核，并在进行fft之前将其填充到图像大小。但这也和前者一样慢，表明瓶颈不在内核的学习阶段，而是在逐点乘法部分。此外，我还看到fourier_conv2d可能仅在内核较大的情况下比conv2d有显著的加速，但它应该至少在时间上与conv2d相当，而不是这么慢。“它应该至少在时间上与conv2d相当”为什么？conv2d每像素执行3x3=9次乘法和加法（MAD）。一次FFT需要的远不止这些。你需要做2次FFT，然后将两个结果相乘。请注意，与将图像数据从RAM获取到CPU并将结果写入RAM相比，9个MAD在现代CPU上花费的时间更少。使用FFT路由，即使只计算了频域的一半，也会有更多的数据，因为内核现在变成了与图像本身一样多的数据。你需要在你的图像上迭代两次，而不是一次。此外，在你链接的论文摘要中，它说“同时多次重复使用相同的转换后的特征图”。也就是说，他们计算所有输入图像的FFT，然后根据转换后的数据训练他们的网络。他们将FFT的成本从等式中剔除。现在它们从每像素9个MAD变为每像素1个乘法。我必须阅读这篇文章才能理解它们对核权重的作用，似乎它们会有更多的核权重。通过在图像上迭代两次，你的意思是一次用于fft，一次用于逐点乘法吗？关于图像/输入数据的fft，在开始训练之前，我已经对整个数据集进行了fft，所以这里不计算时间。对于内核的fft，我也尝试过直接初始化一个complex64内核，而不是初始化真正的内核然后进行fft，所以这个开销也被消除了。但它仍然像这个案子一样慢。所以它让我觉得开销实际上是在逐点乘法部分，也就是复数的乘法。你试过去掉转置吗？您应该能够使用

einsum

而不使用它：

ijkl，ljko->ijko

。虽然

tf.einsum

可能只是速度很慢。