如何在TensorFlow Keras模型中正确实现FFT网络层

• FFT变换轴指定错误：tf.spectral.fft（2.x版本对应tf.signal.fft）默认对输入张量的最后一个维度执行FFT运算。你的测试代码中输入被reshape为(1,64,1)形状，最后一维长度为1，对长度为1的序列做FFT的结果就是输入值本身，虚部始终为0，和你对长度64的完整信号做FFT的预期完全不符。
• 旧版Keras缺失复数类型支持：TensorFlow 1.12.0内置的Keras未原生支持complex64/complex128作为层间传递的标准张量类型，当Lambda层返回复数张量时，Keras内部的输出转换逻辑会自动丢弃虚部、仅保留实部，进一步导致输出看不到虚部结果。同时形状推断逻辑无法识别复数结构，因此模型摘要中输出形状始终显示为(None, None, 1)。
• 对比基准不匹配：你用来做结果参照的np.fft.rfft是实信号专用FFT接口，仅返回非冗余正频率分量，输出长度为输入长度//2 +1；而代码中调用的tf.spectral.fft是通用复数到复数FFT接口，输出长度与输入长度完全一致，二者计算逻辑、输出维度本身存在差异，不能直接对比。

核心思路是显式指定FFT变换轴，同时将复数输出的实部、虚部拆分为实值通道输出，绕开Keras对复数张量的兼容问题，参考代码如下：

import tensorflow as tf
import tensorflow.keras as keras
import numpy as np

# 生成测试信号
s = np.sin(np.linspace(0, 4*np.pi, 64))

# 封装FFT运算逻辑
def fft_ops(v):
    # 压缩最后一维的单通道，得到形状(batch, 信号长度)的实张量
    v = tf.squeeze(v, axis=-1)
    # 转换为复数类型，指定对信号长度维度（轴=1）执行FFT
    v_complex = tf.cast(v, tf.complex64)
    fft_res = tf.spectral.fft(v_complex)
    # 若需要对齐np.fft.rfft结果，替换上一行为 fft_res = tf.spectral.rfft(v_complex)
    # 将实部、虚部堆叠为最后一维的两个通道，输出全实值张量，形状为(batch, 频率点数量, 2)
    return tf.stack([tf.real(fft_res), tf.imag(fft_res)], axis=-1)

# 构建测试模型
inputs = keras.layers.Input(shape=(None, 1))
fft_out = keras.layers.Lambda(fft_ops)(inputs)
model = keras.Model(inputs=inputs, outputs=fft_out)

# 推理验证
y = model.predict(s.reshape(1, 64, 1))
# 拼接回复数格式和numpy FFT结果对比
tf_fft_result = y[0, :, 0] + 1j * y[0, :, 1]
np_fft_result = np.fft.fft(s)
print("TF结果与Numpy结果最大误差：", np.max(np.abs(tf_fft_result - np_fft_result)))

运行后误差在浮点精度范围内，结果完全对齐。如果升级到TensorFlow 2.6及以上版本，Keras已经基础支持复数张量传递，可以不用拆分实虚部，但仍需要显式指定FFT的变换轴，避免出现对最后一维单通道做变换的错误。

• Lambda层封装的FFT操作完全支持误差反向传播。tf.spectral.fft/tf.signal.fft是TensorFlow内置的可微算子，梯度会自动沿FFT逆变换路径回传至前置网络层，不需要额外自定义梯度逻辑。
• 实际使用时建议将所有复数相关运算（比如幅值计算、相位计算、频域加权等）都放在Lambda层内部完成，层间只传递实值张量，避免复数类型引发的兼容问题。

内容的提问来源于stack exchange，提问作者G S