[Python] Fourier Neural Operator in tensorflow nearly good

I try to implement in tensorflow a Fourier Neural Operator from this paper in 2D. As simple example, i have a single layer that learns to perform a gaussian blur. However, there are littles errors especialy on horizontal lines (right image). I dont know why.

Many thanks in advance!

]1 Here is the full code (only few seconds to reach the end).

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from tensorflow.keras import layers, Model, Input
import cv2

class FourierNeuralOperator2D(tf.keras.layers.Layer):
def __init__(self, modes, in_channels, out_channels, **kwargs):
super().__init__(**kwargs)
self.modes = modes
self.in_channels = in_channels
self.out_channels = out_channels

def build(self, input_shape):
self.spectral_weights = self.add_weight(
shape=(self.modes, self.modes, self.in_channels, self.out_channels),
initializer="random_normal",
trainable=True,
name="spectral_weights",
dtype=tf.float32
)

def spectral_convolution(self, x):
x_complex = tf.cast(x, tf.complex64)
x_ft = tf.signal.fft2d(x_complex)
x_ft_shifted = tf.signal.fftshift(x_ft, axes=(1, 2))
H, W = tf.shape(x)[1], tf.shape(x)[2]
start_h = (H - self.modes) // 2
start_w = (W - self.modes) // 2
x_ft_trunc = x_ft_shifted[:, start_h:start_h + self.modes, start_w:start_w + self.modes, :]

spectral_weights_complex = tf.cast(self.spectral_weights, tf.complex64)
out_ft = tf.einsum('bxyi,xyio->bxyo', x_ft_trunc, spectral_weights_complex)

B = tf.shape(x)[0]
padded = tf.zeros([B, H, W, self.out_channels], dtype=tf.complex64)
indices = tf.stack(tf.meshgrid(
tf.range(B),
tf.range(start_h, start_h + self.modes),
tf.range(start_w, start_w + self.modes),
tf.range(self.out_channels),
indexing='ij'
), axis=-1)
padded = tf.tensor_scatter_nd_update(padded, indices, out_ft)

padded_shifted = tf.signal.ifftshift(padded, axes=(1, 2))
output_complex = tf.signal.ifft2d(padded_shifted)
output = tf.math.real(output_complex)
return output
def call(self, inputs):
return self.spectral_convolution(inputs)

def create_blur_model(input_shape=(128, 128, 3)):
inputs = Input(shape=input_shape)
x = FourierNeuralOperator2D(modes=128, in_channels=3, out_channels=3)(inputs)
return Model(inputs=inputs, outputs=x)

x = layers.Conv2D(32, kernel_size=1, padding='same')(inputs)
for i in range(4):
residual = x
x = FourierNeuralOperator2D(modes=64, in_channels=32, out_channels=32)(x)
x = layers.Add()([x, residual])
outputs = layers.Conv2D(3, kernel_size=1, padding='same')(x)
model = Model(inputs=inputs, outputs=outputs)
return model

def create_test_image(height=128, width=128):
img = np.zeros((height, width, 3), dtype=np.float32)
img[30:70, 20:60, 0] = 1.0
y, x = np.ogrid[:height, :width]
center_y, center_x = 60, 100
mask = (x - center_x)**2 + (y - center_y)**2 <= 20**2
img[mask, 1] = 1.0
for i in range(20):
img[90+i, 80-i:80+i+1, 2] = 1.0
return img

def create_gaussian_blur_target(image, sigma=2.0):
blurred = np.zeros_like(image)
for c in range(image.shape[2]):
blurred[:, :, c] = ndimage.gaussian_filter(image[:, :, c], sigma=sigma)
return blurred.astype(np.float32)

def train_blur_model(model):
print("Création des données d'entraînement...")
images = []
targets = []
for i in range(400):
img = create_test_image()
noise = np.random.normal(0, 0.05, img.shape).astype(np.float32)
img = np.clip(img + noise, 0, 1)
sigma = np.random.uniform(1.5, 3.0)
target = create_gaussian_blur_target(img, sigma=sigma)
images.append(img)
targets.append(target)
X_train = np.array(images)
y_train = np.array(targets)
print(f"Données créées: {X_train.shape} -> {y_train.shape}")

model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
loss='mse',
metrics=['mae']
)

print("Entraînement du modèle...")
history = model.fit(
X_train, y_train,
epochs=100,
batch_size=8,
validation_split=0.2,
verbose=1
)
return model, history

def compare_results(model):
test_img = create_test_image()
test_input = np.expand_dims(test_img, 0)
fft_result = model.predict(test_input)[0]
gaussian_result = create_gaussian_blur_target(test_img, sigma=2.0)

fig, axes = plt.subplots(1, 4, figsize=(15, 5))
axes[0].imshow(test_img)
axes[0].set_title('Image originale')
axes[0].axis('off')
axes[1].imshow(np.clip(fft_result, 0, 1))
axes[1].set_title('FNO (appris)')
axes[1].axis('off')
axes[2].imshow(gaussian_result)
axes[2].set_title('Flou gaussien (référence)')
axes[2].axis('off')
axes[3].set_title('Abssolute difference')
axes[3].axis('off')
axes[3].imshow(np.abs(fft_result - gaussian_result))
plt.tight_layout()
plt.show()
out = np.abs(fft_result - gaussian_result)
out = (255*out).astype(np.uint8)
cv2.imwrite('out.png', out)
mse = np.mean((fft_result - gaussian_result)**2)
print(f"Erreur MSE entre FNO et référence: {mse:.6f}")
return fft_result, gaussian_result

def analyze_learned_kernel(model):
fft_layer = model.layers[1]
weights = fft_layer.get_weights()[0]
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for channel in range(3):
ax = axes[channel]
w = np.abs(weights[:, :, channel, channel])
im = ax.imshow(w, cmap='viridis')
ax.set_title(f'Poids spectraux (canal {channel})')
ax.axis('off')
plt.colorbar(im, ax=ax)
plt.tight_layout()
plt.show()

# Exécution
model = create_blur_model(input_shape=(128, 128, 3))
model.summary()
model, history = train_blur_model(model)

print("\n=== Comparaison des résultats ===")
fft_result, gaussian_result = compare_results(model)

print("\n=== Analyse des poids spectraux appris ===")
analyze_learned_kernel(model)

print("\n=== Courbes d'apprentissage ===")
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Val Loss')
plt.title('Evolution de la Loss')
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.legend()
plt.grid(True)
plt.subplot(1, 2, 2)
plt.plot(history.history['mae'], label='Train MAE')
plt.plot(history.history['val_mae'], label='Val MAE')
plt.title('Evolution du MAE')
plt.xlabel('Epoch')
plt.ylabel('Mean Absolute Error')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

print("\n=== Test réussi ! ===")
print("Votre FNO a appris à faire du flou !")

Continue reading...