Skip to content

Instantly share code, notes, and snippets.

@0xfe
Last active Mar 2, 2020
Embed
What would you like to do?
Neural Network to learn the Discrete Fourier Transform
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
batch_size = 15000
fft_size = 128
xs = np.random.randn(batch_size, fft_size) + np.random.randn(batch_size, fft_size) * 1j
ys = np.fft.fft(xs, axis=1)
(Xs, Ys) = (np.stack([xs.real, xs.imag], axis=2), np.hstack([ys.real, ys.imag]))
print(Xs.shape, Ys.shape)
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(fft_size,2)),
tf.keras.layers.Dense(fft_size*2, use_bias=False)
])
optimizer = tf.keras.optimizers.Adam(amsgrad=True)
model.compile(loss='mean_squared_error', optimizer=optimizer)
history = model.fit(Xs, Ys, epochs=100)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment