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

时间:2019-04-01 10:36:59

标签: python tensorflow keras fft conv-neural-network

根据卷积定理,卷积在傅立叶域中变为点向乘法,并且在许多先前的工作中,由于将卷积运算转换为点向乘法运算而获得的增益已显示出进行傅里叶变换的开销被掩盖了类似于以下内容-https://arxiv.org/abs/1312.5851.

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

在该层中,调用功能实现如下: 注意-在我的数据集中,即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)

1 个答案:

答案 0 :(得分:-1)

我不认为您的结果令人惊讶,在Keras中实现Conv2D留给了后端,并且大多数后端(例如TensorFlow)都具有非常优化的卷积操作版本,特别是如果您使用CuDNN。因此,您自己的版本(比天真的实现要快)要比高度优化的版本要慢。

为了进行有意义的比较,您可能必须实现一个基线Conv2D,它可以以天真的方式进行卷积,而无需进行任何优化。