Keras/Tensorflow - conv2d 的傅里叶逐点乘法实现比空间卷积慢 4 倍

Question

根据卷积定理，卷积在傅立叶域中变为逐点乘法，并且由于在许多先前的工作中将卷积运算转换为逐点乘法运算，已证明进行傅立叶变换的开销被增益所掩盖像下面这样 - https://arxiv.org/abs/1312.5851.

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

在层中，调用函数实现如下—— 注意 - 在我的数据集中，即 MNIST 图像高度 = 宽度，因此转置工作正常

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 层给出的精度，但它的运行速度比 Conv2D 慢了大约 4 倍，因此转换为傅立叶域的目的就落空了。谁能解释一下为什么会发生这种情况以及如何在傅立叶域中复制快速卷积的结果？

（注意 - 对于任何可能觉得 tf.spectral.rfft2d(self.kernel) 可能是开销的人来说，事实并非如此。

另外，我认为 Conv2D 函数可能会展平 4D 输入张量和内核，以将其减少为矩阵乘法，如此处所述 - enter link description here。除了像我对 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)

Answer 1

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

为了进行有意义的比较，您可能必须实现一个基线 Conv2D，它以一种简单的方式进行卷积，没有任何类型的优化。