How to implement a Fourier Convolution layer in keras?

Question

I'm currently investigating the paper FCNN: Fourier Convolutional Neural Networks. The main contribution of the paper is that CNN training is entirely shifted to the Fourier domain without loss of effectiveness. The proposed architecture looks as follows:

The authors state that the implementation was done in keras, however, it is not publicly available. I know I can define a Fourier transformation in the following way:

model.add(layers.Lambda(lambda v: tf.real(tf.spectral.rfft(v))))

But this is not a Fourier Convolution, right? How should I go from here?

Answer 1

An FFT-based convolution can be broken up into 3 parts: an FFT of the input images and the filters, a bunch of element-wise products followed by a sum across input channels, and then an IFFT of the outputs (Source).

Or as it is written in the paper:

So, for a Fourier Convolution Layer you need to:

1. Take the input layer and transform it to the Fourier domain: input_fft = tf.spectral.rfft2d(input)

2. Take each kernel and transform it to the Fourier domain: weights_fft = tf.spectral.rfft2d(layer.get_weights()) Note: The Fourier domain "images" for the input and the kernels need to be of the same size.

3. Perform element-wise multiplication between the input's Fourier transform and Fourier transform of each of the kernels: conv_fft = keras.layers.Multiply(input_fft, weights_fft)

4. Perform Inverse Fourier transformation to go back to the spatial domain. layer_output = tf.spectral.irfft2d(conv_fft)

Note: I used pseudo-code, it will probably needs some tuning for it to actually to work.