做信号处理的人应该都经历过这种事：手里一堆波形数据，想喂进 TensorFlow 之前先做快速傅里叶变换（FFT）；等数据量大到单机预处理扛不住，又开始琢磨能不能在分布式训练流程里顺手把频谱算出来。今天这篇就专门聊聊 TensorFlow 里分布式 FFT 的正确打开方式。我会把 API 细节、数据管道怎么搭、多机多卡怎么配，以及我实际跑过之后觉得最值得注意的坑都写出来。适合正在做语音、振动、地震、雷达这些频域特征相关训练的朋友参考。以前我总习惯先把波形离线转成频谱，存成 npy 或者 TFRecord，再去做训练。样本量一上来，这条路越来越难受：离线转换要跑好几个小时，想换个窗函数或者 fft_length 又得全量重跑，分布式训练的时候还要先找一台机器把数据处理好再分发，绕远路不说，还容易出错。后来我把 FFT 从“离线预处理”挪进了 TensorFlow 的数据管道和模型入口，配合tf.distribute做多机多卡训练，整个流程才算顺起来。当然，代价是实打实的：FFT 输出是复数，绝大多数模型层吃不了复数；分布式数据分片、批次大小、fft_length 这些细节每个都可能让你调试一晚上。1. 为什么我放弃了“先离线存频谱，再训练”的老路1.1 离线预处理在规模化后的三个硬伤先说存储。一段 1024 点 float32 波形，原始大小是 4KB 左右。用rfft转成复数频谱，输出 513 个 complex64，每个 8 字节，约 4KB，看上去和原始波形差不多。但如果你做的是短时傅里叶变换（STFT），情况就完全不同了：一段 8192 点的波形，frame_length=256、frame_step=128，会产生约 63 帧，每帧 129 个复数频点，光频谱数据就是 63 × 129 × 8 字节，接近 65KB，比原始波形放大了 16 倍。大规模数据集这么一搞，存储直接翻一个数量级。再说 epoch 循环。训练的时候每个 epoch 都要重新读一遍频谱数据，这本不稀奇。但如果你想在频域做数据增广，比如时间偏移、加噪、频点掩码，离线把频谱固定下来之后，这些都玩不了。更麻烦的是参数调试：今天觉得 fft_length 用 2048 效果不好，明天想加一个 Hann 窗，改了参数就得从头把离线预处理任务重跑一遍。在多机环境下，这意味着还要维护一个独立的预处理程序和一个训练程序，两边容易版本不同步，出错概率成倍增加。1.2 分布式训练流程里的 FFT，本质上是把“预处理”变成“流水线算子”TensorFlow 的训练框架天然是数据并行的。FFT 作为一个逐样本或逐 batch 的确定性算子，完全可以放在数据读取之后、模型第一步之前。每个 worker 读取自己分片的数据，在tf.data管道里完成 FFT，再喂给模型。这样变换参数只改一个配置，不用重新跑离线任务；数据增广也能自由调整。但这句话听着简单，做起来细节很多。FFT 输出是复数，Keras 里大部分层不支持 complex dtype；分布式策略对数据管道的切分有讲究；GPU 和 CPU 上 FFT 的编排方式不一样。下面逐个拆开讲。2. tf.signal 里的 FFT 全家桶：API 细节与 NumPy 的核心差异2.1 常用算子速查TensorFlow 把 FFT 相关操作都放在tf.signal命名空间下，我从项目里实际用到的整理了一张表：算子输入类型输出类型作用维度tf.signal.fftcomplex64/complex128同输入最后一维tf.signal.ifftcomplex64/complex128同输入最后一维，逆变换tf.signal.rfftfloat32/float64complex64/complex128最后一维，实数输入tf.signal.irfftcomplex64/complex128float32/float64最后一维，逆变换tf.signal.fft2dcomplex64/complex128同输入最后两维tf.signal.rfft2dfloat32/float64complex64/complex128最后两维tf.signal.fft3dcomplex64/complex128同输入最后三维tf.signal.rfft3dfloat32/float64complex64/complex128最后三维tf.signal.stftfloat32/float64complex64最后一维上分帧加窗变换核心区别在于，这些算子是“固定维度”的，不能像np.fft.fft(a, axis=0)那样指定作用轴。tf.signal.fft永远只对张量的最后一维做一维变换，fft2d永远只对最后两维做变换。这是和 NumPy 习惯最大的不同，也是新手最容易踩的坑。2.2 最容易踩的“axis”坑举个具体例子。假设我有一条振动波形数据，shape 是[batch, timesteps, channels]，想沿着timesteps轴做 FFT。因为 timesteps 在中间而不是最后一维，tf.signal.rfft会把它当成“batch × timesteps × channels”里的最后一维去处理，也就是对 channels 做变换，结果完全错误。正确做法是先把时间维挪到最后一维：# x: [batch, timesteps, channels] -> [batch, channels, timesteps] x = tf.transpose(x, [0, 2, 1]) spec = tf.signal.rfft(x) # 现在按时间轴做 FFT x = tf.transpose(spec, [0, 2, 1]) # 如果需要再换回来所以我的经验是：在设计输入数据布局时，尽量让需要变换的那一维落在最后一个维度，否则整个预处理会多一堆 transpose，还容易搞混。这个问题在 2D 图像数据上更明显，后面单独讲。2.3 为什么我坚持用 rfft 而不是 fft真实采集的波形数据几乎都是实数，而实数信号的频谱具有共轭对称性，正负频率部分信息重复。tf.signal.rfft只输出一半的频点，比如输入 1024 点，输出 513 个复数点。只保留一半，意味着计算量、内存和后续通信量都几乎减半。别小看这 2 倍。在单机上可能无所谓，但在多机分布式训练里，这个“减半”作用在每批次数据、每个 worker 上都生效。我们做大规模振动数据训练时，batch 内所有样本的频谱都要从 CPU 或 GPU 传到显存里，省一半就是实打实的省一半显存带宽。另外，如果后续要落盘缓存频谱，complex64 的存储也能比 complex128 省一半。3. 把 FFT 放到哪一个阶段？数据管道和模型层的实战取舍3.1 方案 A：在 tf.data 管道内做 FFT（我默认推荐的方案）把 FFT 放在数据管道里，是分布式场景下最稳的做法。它的本质是把 FFT 当成普通的数据预处理算子，由每个 worker 的map并发执行。我的标准写法是这样：def rfft_to_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) # 确保 float32 spectrum = tf.signal.rfft(waveform, fft_length=[8192]) feature = tf.stack( [tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1 ) # [4097, 2] return feature, label dataset = tf.data.Dataset.from_tensor_slices((waveforms, labels)) dataset = dataset.shuffle(10000) dataset = dataset.map(rfft_to_feature, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.batch(64) dataset = dataset.prefetch(tf.data.AUTOTUNE)这里有几个顺序性细节，是我反复改过才固定的：第一，shuffle要在map之前。如果先做 FFT 再 shuffle，相当于把每个样本的频谱先算出来再打乱，FFT 计算量白白浪费，还占内存。尤其做 STFT 时，先 map 后 shuffle 会让 shuffle 缓冲区里的内存开销爆炸，因为每个样本从一维波形变成了二维频谱图。第二，rfft的输出是复数，如果在 map 里不拆实部虚部，后面batch出来的张量是 complex64，Keras 模型基本没法直接用。第三，拆实部虚部的轴要选对：tf.stack([real, imag], axis=-1)会把每个频点变成一个长度为 2 的向量，最终形状是[num_freq_bins, 2]。如果你希望模型把实虚当成两个分离通道，也可以改用axis=0或tf.concat，按需调整。3.2 方案 B：把 FFT 写成 Keras 自定义层放进模型另一种做法是把 FFT 写成一个自定义层，放在模型的入口。典型代码如下：class FFTFeatureLayer(tf.keras.layers.Layer): def __init__(self, fft_length, **kwargs): super().__init__(**kwargs) self.fft_length = fft_length def call(self, inputs): inputs = tf.cast(inputs, tf.float32) spectrum = tf.signal.rfft(inputs, fft_length=[self.fft_length]) return tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) def compute_output_shape(self, input_shape): length = self.fft_length if self.fft_length is not None else input_shape[-1] return input_shape[:-1] + (length // 2 + 1, 2)注意，这个层必须在内部就把复数拆成实数特征返回，不能把 complex64 直接暴露给后面的 Conv1D、Dense 这些层。Keras 的自动微分和优化器对复数 dtype 的支持不完整，强行往下传很容易出现不可预测的报错。方案 B 的优势在于，FFT 计算发生在训练设备上。如果你的样本很长，FFT 在 GPU 上有 cuFFT 加速，会比在 CPU 数据管道里快不少。另外，FFT 是线性且可微的算子，放进模型内部后，梯度可以流回频谱甚至流回原始波形，这就为后面要讲的“可学习频域滤波”留了可能性。3.3 怎么选：一张对比表和我的经验维度数据管道内 FFT自定义模型层 FFT执行位置worker 的 CPU/IO 线程训练设备（GPU/TPU）多机分布式适配天然按数据分片，基本不用额外处理跟随模型 replica，复杂度和序列化风险更高模型结构模型保持干净，输入直接是特征输入到模型入口是原始波形，第一层做变换反向传播参与性无法参与训练，只能做固定预处理可以参与端到端学习，做可微频谱处理调试友好度出了问题，管道算子独立排查需要连带模型一起排查批量离线 FFT可以通过 model.predict 完成同样能通过 model.predict 完成我的实际经验是：如果目标是分类、回归这类常规任务，FFT 只是一个固定的特征提取步骤，那就放数据管道，简单、稳、好排查。只有当任务的核心诉求是“在频谱域做可学习处理”时，才把 FFT 放进模型内部，因为这时候你需要的是把梯度从频域传回时域的能力。4. 多机多卡实战：数据分片、TFRecord 与 MultiWorkerMirroredStrategy4.1 输入数据存什么：时域还是频谱？我的默认建议是 TFRecord 里存原始时域波形，把 FFT 留给训练管道。这样做的最大好处是灵活：想改 fft_length、改窗函数、改 STFT 参数，不用重新生成数据集。缺点是在训练时需要额外计算 FFT，IO 带宽也可能成为瓶颈。反过来，如果任务已经定型，STFT 特征非常大，或者共享存储的 IO 实在紧张，也可以离线一次性把频谱算好，存成 TFRecord，训练时直接读频谱。代价是以后每次改参数都要重新生成一遍数据。这个取舍要看项目阶段：早期探索阶段一定要存时域；模型结构定型、进入大规模重复训练阶段，再考虑缓存频谱。4.2 distribute_datasets_from_function 的正确打开方式多 worker 训练时，最容易犯的错误是把数据集直接strategy.experimental_distribute_dataset(dataset)，然后在每个 worker 上整个数据集都重复载入。更灵活、也更高效的方式是用distribute_datasets_from_function，在每个 worker 上按input_context做分片，各读各的文件，各自做 FFT。我整理过一个完整可跑的骨架，核心是文件级分片 + 管道内 FFT：import tensorflow as tf GLOBAL_BATCH_SIZE = 64 FILE_PATTERN = "/shared/data/train/*.tfrecord" def parse_tfrecord(example_proto): feature_description = { "waveform": tf.io.FixedLenFeature([8192], tf.float32), "label": tf.io.FixedLenFeature([], tf.int64), } parsed = tf.io.parse_single_example(example_proto, feature_description) return parsed["waveform"], parsed["label"] def rfft_to_feature(waveform, label): spectrum = tf.signal.rfft(waveform, fft_length=[4096]) feature = tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) return feature, label def make_distributed_dataset(input_context, file_pattern): # 每个 worker 只读属于自己的那部分文件 files = tf.data.Dataset.list_files(file_pattern, shuffle=True) files = files.shard( input_context.num_input_pipelines, input_context.input_pipeline_id ) ds = files.interleave( tf.data.TFRecordDataset, num_parallel_calls=tf.data.AUTOTUNE, deterministic=False, ) ds = ds.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE) ds = ds.shuffle(10000) ds = ds.map(rfft_to_feature, num_parallel_calls=tf.data.AUTOTUNE) per_replica_batch = input_context.get_per_replica_batch_size(GLOBAL_BATCH_SIZE) ds = ds.batch(per_replica_batch, drop_remainder=True) ds = ds.prefetch(tf.data.AUTOTUNE) return ds strategy = tf.distribute.MultiWorkerMirroredStrategy() train_ds = strategy.distribute_datasets_from_function( lambda input_context: make_distributed_dataset(input_context, FILE_PATTERN) ) with strategy.scope(): model = tf.keras.Sequential([ tf.keras.layers.Input(shape=(2049, 2)), # 4096//2+1 个频点，每个频点实部+虚部 tf.keras.layers.Flatten(), tf.keras.layers.Dense(64, activation="relu"), tf.keras.layers.Dense(2, activation="softmax"), ]) model.compile( optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"], ) model.fit(train_ds, epochs=10)这里面最容易被忽略的是input_context.get_per_replica_batch_size(GLOBAL_BATCH_SIZE)。它会把全局 batch size 按照 worker 数量切成每个 worker 自己应处理的副本大小。如果你直接写batch(64)，在多 worker 下每个 worker 都读 64 条，总 batch 就不对了，梯度的 all-reduce 语义也会出问题。另外，文件级分片要求文件数量远大于 worker 数量。如果只有一个大文件，或者文件数量小于 worker 数，会出现某些 worker 分不到数据。这种场景可以改成记录级分片：ds = tf.data.TFRecordDataset(files) ds = ds.shard( input_context.num_input_pipelines, input_context.input_pipeline_id )记录级分片会打散到单条记录，但 IO 性能略差一些，因为每个 worker 都要打开同一批文件去读取自己需要的记录。文件级分片仍然是优先选择。4.3 多 worker 下批次大小、复数和标签对齐的细节先说标签。很多第一次写的人会误以为频谱维度变了，标签也要跟着变。其实标签是对“整条样本”而言的，一个 waveform 对应一个 label，不管 rfft 之后频点数变成多少，label 始终保持不变，注意别在 map 里把 label 的维度不小心压掉就行。再看模型输入。原始波形 8192 点，fft_length=[4096]表示只取前 4096 点参与 FFT，输出 2049 个复数频点，每个频点拆成实部和虚部，最终特征 shape 是(2049, 2)。对应到模型Input(shape=(2049, 2))，这两维千万不要搞反。如果你用的是tf.signal.stft，输出变成(frames, num_freq_bins)，模型输入还要相应调整。最后是复数问题。rfft 输出 complex64，拆成实部和虚部之后变成两个实数张量。手写tf.stack是在管道里处理；如果放在模型层里，同样要拆，否则后面接 Conv2D 或者 Dense 会直接报 dtype 不支持。这个坑我踩了不止一次，尤其从单卡切到多卡时，踩的频率更高。4.4 不训练模型，只想用 TensorFlow 做批量分布式 FFT有一个容易忽视的用法：如果你手头有海量波形需要转成频谱，但暂时不训练模型，完全可以借用 TensorFlow 的分布式能力来跑批量 FFT，没必要再单独起一个大数据工具链。做法很简单：把前面的FFTFeatureLayer包进一个Sequential模型，编译之后直接model.predict(dataset)。prediction 阶段照样走distribute_datasets_from_function做数据分片，每个 worker 在各自的副本上执行 FFT，最后把结果拼回来。这个技巧我在做某次批量频谱转换时用过，省掉了专门的分布式预处理程序。5. 跑起来之后才会注意到的性能细节5.1 num_parallel_calls 不是越大越好tf.data里的map有num_parallel_calls参数，很多人一看并行就调到很大。实测下来并非如此。FFT 本身是计算密集型操作，在 CPU 上跑会占满核心；如果你的num_parallel_calls大于物理核数，线程切换的开销反而会让吞吐下降。我做过一个对比：同一个rfft_to_feature管道，在 8 核机器上把num_parallel_calls从 4 调到 16，每秒处理样本数不升反降，大约降到原来的 70% 左右。我的建议是直接用tf.data.AUTOTUNE，让框架自己根据机器负载去调节。如果非要手动设，先看 CPU 核心数，设置为物理核心数附近，不要盲目拉满。5.2 fft_length、窗函数和频谱泄漏fft_length这个参数很容易被忽视，但它直接影响频谱形态。如果输入长度大于fft_length，tf.signal.rfft会隐式截断，直接丢掉后面多出来的点；如果输入长度小于fft_length，会隐式补零。很多人以为补零只是把点数变多、看起来更平滑，实际上补零并不会增加真实频率分辨率，它只是在连续频谱上做了更多采样，视觉上变平滑而已。更隐蔽的是截断：如果你怀疑高频段丢失了信息，先检查是不是输入长度和fft_length不匹配。靠窗函数减少频谱泄漏也是必修课。实际采集波形往往不是整周期截断，直接 FFT 会导致频谱泄漏，表现为峰附近出现额外旁瓣。在管道里加窗很简单：WINDOW = tf.signal.hann_window(4096, dtype=tf.float32) def rfft_to_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) waveform = waveform[:4096] * WINDOW # 先手动截断，再加窗 spectrum = tf.signal.rfft(waveform, fft_length=[4096]) feature = tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) return feature, label注意，窗口长度必须和真正参与 FFT 的点数一致。如果你直接对 8192 点波形加 4096 点的窗，然后期望fft_length=[4096]截断后正好和窗口匹配，这是错的。因为 rfft 是在加窗之后的整个序列上做的，内部先截断还是先加窗，行为和你手写代码时的预期可能不一样。我自己通常是先显式切片到指定长度，再乘窗口，最后才做 rfft，这样每一步都可控。5.3 混合精度与 dtype 问题训练里开了自动混合精度之后，很多层会自动用 float16 计算。但tf.signal.rfft对输入 dtype 有要求，通常只支持 float32 或 float64。如果你在数据管道早期就把波形 cast 成 float16，后面调用 rfft 时要么报错，要么需要框架隐式转换，既慢又损失精度。我的习惯是在管道里统一保证波形是 float32，rfft 输出自然是 complex64。等到频谱拆成实部虚部之后，模型内部是否转 float16 由混合精度策略去决定，不要把 dtype 问题带到 FFT 算子里。5.4 多 worker 下的 IO 与缓存多 worker 训练时，如果每个 epoch 都重新从共享存储读原始波形，IO 压力会非常大。一个折中方案是dataset.cache()：把经过 parse 和初步清洗的数据缓存到本地磁盘，避免每一轮都重复读共享存储。但缓存要慎用，如果数据量大到几个 TB，本地磁盘放不下，缓存就失去意义。文件级分片对 IO 的优化比缓存更直接。因为它避免多个 worker 重复读同一份文件。如果你的文件数量足够多，文件级分片 +interleave+prefetch通常就能把 IO 压到一个可控水平。如果 IO 依然紧张，再考虑离线频谱化。6. 进阶玩法：STFT、二维 FFT 和可学习的频域滤波6.1 STFT：把一维波形变成时间-频率图语音、振动这一类非平稳信号，单次 FFT 会把整个时间段的频率特征混在一起，丢失时间信息。更常见的做法是做 STFT，得到一张时间-频率图，再交给 CNN 去分类。TensorFlow 里tf.signal.stft把分帧、加窗、FFT 合在了一起：def waveform_to_stft_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) spectrogram = tf.signal.stft( waveform, frame_length=256, frame_step=128, fft_length=256, window_fn=tf.signal.hann_window, pad_end=True, ) # 输出 [frames, 129] complex64 spectrogram = tf.abs(spectrogram) spectrogram = tf.math.log1p(spectrogram) # 压缩动态范围 return spectrogram, labelframe_length是窗长，frame_step是帧移，fft_length参与 FFT 的点数。一般建议fft_length等于frame_length，这样窗口函数和 FFT 截断区域完全一致。pad_end=True会在信号尾部补零，保证最后一帧能被处理，避免丢掉尾部信息。STFT 生成的特征比一维频谱大得多。在分布式场景里，如果每个 worker 的 CPU 管道里都跑 STFT，计算量可不小。我的经验是：单样本帧数不多、数据量中等时，放数据管道没问题；如果帧数很多，或者每帧频点很大，可以考虑把 STFT 写进自定义层，让 GPU 去并行。6.2 二维 FFT：图像和雷达数据的维度游戏图像或者雷达距离-多普勒图这类数据，二维 FFT 更常用。tf.signal.rfft2d只作用于最后两维，所以灰度图像[batch, height, width]可以直接用：spectrum = tf.signal.rfft2d(image, fft_length=[height, width]) # 输出 [batch, height, width//2+1] complex64但如果图像是[batch, height, width, channels]，想对每个通道分别做二维 FFT，绝对不能用原始形状直接丢进去，否则rfft2d会对width和channels做变换，结果完全不对。需要先把通道维挪到前面：image_t = tf.transpose(image, [0, 3, 1, 2]) # [batch, channels, height, width] spectrum = tf.signal.rfft2d(image_t) # 对 height/width 做 FFT这个问题本质上还是之前说的“固定作用维度”问题。凡是遇到多维输入，先停下来想清楚：要变换的轴是不是最后两维？如果不是，先 transpose 再操作。6.3 在频谱域做可学习滤波，让 FFT 参与反向传播FFT 是线性算子，TensorFlow 对它也做了梯度支持。利用这一点，可以在模型中间插入一个“频域滤波层”：先把波形 FFT 到频域，乘一个可学习的 mask，再 IFFT 回时域，让模型自己学会在频域上怎么滤波。这在语音增强、去噪这类任务里非常实用。简化实现大致是这样：class LearnableFreqMaskLayer(tf.keras.layers.Layer): def __init__(self, fft_length): super().__init__() self.fft_length = fft_length def build(self, input_shape): self.mask = self.add_weight( name="freq_mask", shape=(self.fft_length // 2 + 1,), initializer="ones", trainable=True, ) def call(self, inputs): spectrum = tf.signal.rfft(inputs, fft_length=[self.fft_length]) spectrum = spectrum * tf.cast(self.mask, spectrum.dtype) return tf.signal.irfft(spectrum, fft_length=[self.fft_length])这个 mask 是模型的一个全局可训练变量，分布式训练时它会参与梯度同步，所以多卡也自然适用。这类层放在模型中间，前面提到的方案 A 就无能为力了，必须用自定义层。这也是我说“先看任务需求再决定 FFT 放哪一层”的原因。最后聊一个实践中的小判断标准。我一般默认把 FFT 放进数据管道，只有当任务本身要求“在频谱上做可学习处理”，也就是需要把频谱的梯度传回原始波形时，才把它放进模型内部。数据管道版本的好处是简单、稳、好排查；模型内部版本的好处是灵活、能端到端优化，但代价是你必须把复数谱拆好，还要面对模型序列化、多 worker 下发等一堆细节。两种方案我都跑过，如果你刚开始搭，建议先从数据管道版本开始，跑通之后再决定要不要往模型里挪。
做信号处理的人应该都经历过这种事：手里一堆波形数据，想喂进 TensorFlow 之前先做快速傅里叶变换（FFT）；等数据量大到单机预处理扛不住，又开始琢磨能不能在分布式训练流程里顺手把频谱算出来。今天这篇就专门聊聊 TensorFlow 里分布式 FFT 的正确打开方式。我会把 API 细节、数据管道怎么搭、多机多卡怎么配，以及我实际跑过之后觉得最值得注意的坑都写出来。适合正在做语音、振动、地震、雷达这些频域特征相关训练的朋友参考。以前我总习惯先把波形离线转成频谱，存成 npy 或者 TFRecord，再去做训练。样本量一上来，这条路越来越难受：离线转换要跑好几个小时，想换个窗函数或者 fft_length 又得全量重跑，分布式训练的时候还要先找一台机器把数据处理好再分发，绕远路不说，还容易出错。后来我把 FFT 从“离线预处理”挪进了 TensorFlow 的数据管道和模型入口，配合tf.distribute做多机多卡训练，整个流程才算顺起来。当然，代价是实打实的：FFT 输出是复数，绝大多数模型层吃不了复数；分布式数据分片、批次大小、fft_length 这些细节每个都可能让你调试一晚上。1. 为什么我放弃了“先离线存频谱，再训练”的老路1.1 离线预处理在规模化后的三个硬伤先说存储。一段 1024 点 float32 波形，原始大小是 4KB 左右。用rfft转成复数频谱，输出 513 个 complex64，每个 8 字节，约 4KB，看上去和原始波形差不多。但如果你做的是短时傅里叶变换（STFT），情况就完全不同了：一段 8192 点的波形，frame_length=256、frame_step=128，会产生约 63 帧，每帧 129 个复数频点，光频谱数据就是 63 × 129 × 8 字节，接近 65KB，比原始波形放大了 16 倍。大规模数据集这么一搞，存储直接翻一个数量级。再说 epoch 循环。训练的时候每个 epoch 都要重新读一遍频谱数据，这本不稀奇。但如果你想在频域做数据增广，比如时间偏移、加噪、频点掩码，离线把频谱固定下来之后，这些都玩不了。更麻烦的是参数调试：今天觉得 fft_length 用 2048 效果不好，明天想加一个 Hann 窗，改了参数就得从头把离线预处理任务重跑一遍。在多机环境下，这意味着还要维护一个独立的预处理程序和一个训练程序，两边容易版本不同步，出错概率成倍增加。1.2 分布式训练流程里的 FFT，本质上是把“预处理”变成“流水线算子”TensorFlow 的训练框架天然是数据并行的。FFT 作为一个逐样本或逐 batch 的确定性算子，完全可以放在数据读取之后、模型第一步之前。每个 worker 读取自己分片的数据，在tf.data管道里完成 FFT，再喂给模型。这样变换参数只改一个配置，不用重新跑离线任务；数据增广也能自由调整。但这句话听着简单，做起来细节很多。FFT 输出是复数，Keras 里大部分层不支持 complex dtype；分布式策略对数据管道的切分有讲究；GPU 和 CPU 上 FFT 的编排方式不一样。下面逐个拆开讲。2. tf.signal 里的 FFT 全家桶：API 细节与 NumPy 的核心差异2.1 常用算子速查TensorFlow 把 FFT 相关操作都放在tf.signal命名空间下，我从项目里实际用到的整理了一张表：算子输入类型输出类型作用维度tf.signal.fftcomplex64/complex128同输入最后一维tf.signal.ifftcomplex64/complex128同输入最后一维，逆变换tf.signal.rfftfloat32/float64complex64/complex128最后一维，实数输入tf.signal.irfftcomplex64/complex128float32/float64最后一维，逆变换tf.signal.fft2dcomplex64/complex128同输入最后两维tf.signal.rfft2dfloat32/float64complex64/complex128最后两维tf.signal.fft3dcomplex64/complex128同输入最后三维tf.signal.rfft3dfloat32/float64complex64/complex128最后三维tf.signal.stftfloat32/float64complex64最后一维上分帧加窗变换核心区别在于，这些算子是“固定维度”的，不能像np.fft.fft(a, axis=0)那样指定作用轴。tf.signal.fft永远只对张量的最后一维做一维变换，fft2d永远只对最后两维做变换。这是和 NumPy 习惯最大的不同，也是新手最容易踩的坑。2.2 最容易踩的“axis”坑举个具体例子。假设我有一条振动波形数据，shape 是[batch, timesteps, channels]，想沿着timesteps轴做 FFT。因为 timesteps 在中间而不是最后一维，tf.signal.rfft会把它当成“batch × timesteps × channels”里的最后一维去处理，也就是对 channels 做变换，结果完全错误。正确做法是先把时间维挪到最后一维：# x: [batch, timesteps, channels] -> [batch, channels, timesteps] x = tf.transpose(x, [0, 2, 1]) spec = tf.signal.rfft(x) # 现在按时间轴做 FFT x = tf.transpose(spec, [0, 2, 1]) # 如果需要再换回来所以我的经验是：在设计输入数据布局时，尽量让需要变换的那一维落在最后一个维度，否则整个预处理会多一堆 transpose，还容易搞混。这个问题在 2D 图像数据上更明显，后面单独讲。2.3 为什么我坚持用 rfft 而不是 fft真实采集的波形数据几乎都是实数，而实数信号的频谱具有共轭对称性，正负频率部分信息重复。tf.signal.rfft只输出一半的频点，比如输入 1024 点，输出 513 个复数点。只保留一半，意味着计算量、内存和后续通信量都几乎减半。别小看这 2 倍。在单机上可能无所谓，但在多机分布式训练里，这个“减半”作用在每批次数据、每个 worker 上都生效。我们做大规模振动数据训练时，batch 内所有样本的频谱都要从 CPU 或 GPU 传到显存里，省一半就是实打实的省一半显存带宽。另外，如果后续要落盘缓存频谱，complex64 的存储也能比 complex128 省一半。3. 把 FFT 放到哪一个阶段？数据管道和模型层的实战取舍3.1 方案 A：在 tf.data 管道内做 FFT（我默认推荐的方案）把 FFT 放在数据管道里，是分布式场景下最稳的做法。它的本质是把 FFT 当成普通的数据预处理算子，由每个 worker 的map并发执行。我的标准写法是这样：def rfft_to_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) # 确保 float32 spectrum = tf.signal.rfft(waveform, fft_length=[8192]) feature = tf.stack( [tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1 ) # [4097, 2] return feature, label dataset = tf.data.Dataset.from_tensor_slices((waveforms, labels)) dataset = dataset.shuffle(10000) dataset = dataset.map(rfft_to_feature, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.batch(64) dataset = dataset.prefetch(tf.data.AUTOTUNE)这里有几个顺序性细节，是我反复改过才固定的：第一，shuffle要在map之前。如果先做 FFT 再 shuffle，相当于把每个样本的频谱先算出来再打乱，FFT 计算量白白浪费，还占内存。尤其做 STFT 时，先 map 后 shuffle 会让 shuffle 缓冲区里的内存开销爆炸，因为每个样本从一维波形变成了二维频谱图。第二，rfft的输出是复数，如果在 map 里不拆实部虚部，后面batch出来的张量是 complex64，Keras 模型基本没法直接用。第三，拆实部虚部的轴要选对：tf.stack([real, imag], axis=-1)会把每个频点变成一个长度为 2 的向量，最终形状是[num_freq_bins, 2]。如果你希望模型把实虚当成两个分离通道，也可以改用axis=0或tf.concat，按需调整。3.2 方案 B：把 FFT 写成 Keras 自定义层放进模型另一种做法是把 FFT 写成一个自定义层，放在模型的入口。典型代码如下：class FFTFeatureLayer(tf.keras.layers.Layer): def __init__(self, fft_length, **kwargs): super().__init__(**kwargs) self.fft_length = fft_length def call(self, inputs): inputs = tf.cast(inputs, tf.float32) spectrum = tf.signal.rfft(inputs, fft_length=[self.fft_length]) return tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) def compute_output_shape(self, input_shape): length = self.fft_length if self.fft_length is not None else input_shape[-1] return input_shape[:-1] + (length // 2 + 1, 2)注意，这个层必须在内部就把复数拆成实数特征返回，不能把 complex64 直接暴露给后面的 Conv1D、Dense 这些层。Keras 的自动微分和优化器对复数 dtype 的支持不完整，强行往下传很容易出现不可预测的报错。方案 B 的优势在于，FFT 计算发生在训练设备上。如果你的样本很长，FFT 在 GPU 上有 cuFFT 加速，会比在 CPU 数据管道里快不少。另外，FFT 是线性且可微的算子，放进模型内部后，梯度可以流回频谱甚至流回原始波形，这就为后面要讲的“可学习频域滤波”留了可能性。3.3 怎么选：一张对比表和我的经验维度数据管道内 FFT自定义模型层 FFT执行位置worker 的 CPU/IO 线程训练设备（GPU/TPU）多机分布式适配天然按数据分片，基本不用额外处理跟随模型 replica，复杂度和序列化风险更高模型结构模型保持干净，输入直接是特征输入到模型入口是原始波形，第一层做变换反向传播参与性无法参与训练，只能做固定预处理可以参与端到端学习，做可微频谱处理调试友好度出了问题，管道算子独立排查需要连带模型一起排查批量离线 FFT可以通过 model.predict 完成同样能通过 model.predict 完成我的实际经验是：如果目标是分类、回归这类常规任务，FFT 只是一个固定的特征提取步骤，那就放数据管道，简单、稳、好排查。只有当任务的核心诉求是“在频谱域做可学习处理”时，才把 FFT 放进模型内部，因为这时候你需要的是把梯度从频域传回时域的能力。4. 多机多卡实战：数据分片、TFRecord 与 MultiWorkerMirroredStrategy4.1 输入数据存什么：时域还是频谱？我的默认建议是 TFRecord 里存原始时域波形，把 FFT 留给训练管道。这样做的最大好处是灵活：想改 fft_length、改窗函数、改 STFT 参数，不用重新生成数据集。缺点是在训练时需要额外计算 FFT，IO 带宽也可能成为瓶颈。反过来，如果任务已经定型，STFT 特征非常大，或者共享存储的 IO 实在紧张，也可以离线一次性把频谱算好，存成 TFRecord，训练时直接读频谱。代价是以后每次改参数都要重新生成一遍数据。这个取舍要看项目阶段：早期探索阶段一定要存时域；模型结构定型、进入大规模重复训练阶段，再考虑缓存频谱。4.2 distribute_datasets_from_function 的正确打开方式多 worker 训练时，最容易犯的错误是把数据集直接strategy.experimental_distribute_dataset(dataset)，然后在每个 worker 上整个数据集都重复载入。更灵活、也更高效的方式是用distribute_datasets_from_function，在每个 worker 上按input_context做分片，各读各的文件，各自做 FFT。我整理过一个完整可跑的骨架，核心是文件级分片 + 管道内 FFT：import tensorflow as tf GLOBAL_BATCH_SIZE = 64 FILE_PATTERN = "/shared/data/train/*.tfrecord" def parse_tfrecord(example_proto): feature_description = { "waveform": tf.io.FixedLenFeature([8192], tf.float32), "label": tf.io.FixedLenFeature([], tf.int64), } parsed = tf.io.parse_single_example(example_proto, feature_description) return parsed["waveform"], parsed["label"] def rfft_to_feature(waveform, label): spectrum = tf.signal.rfft(waveform, fft_length=[4096]) feature = tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) return feature, label def make_distributed_dataset(input_context, file_pattern): # 每个 worker 只读属于自己的那部分文件 files = tf.data.Dataset.list_files(file_pattern, shuffle=True) files = files.shard( input_context.num_input_pipelines, input_context.input_pipeline_id ) ds = files.interleave( tf.data.TFRecordDataset, num_parallel_calls=tf.data.AUTOTUNE, deterministic=False, ) ds = ds.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE) ds = ds.shuffle(10000) ds = ds.map(rfft_to_feature, num_parallel_calls=tf.data.AUTOTUNE) per_replica_batch = input_context.get_per_replica_batch_size(GLOBAL_BATCH_SIZE) ds = ds.batch(per_replica_batch, drop_remainder=True) ds = ds.prefetch(tf.data.AUTOTUNE) return ds strategy = tf.distribute.MultiWorkerMirroredStrategy() train_ds = strategy.distribute_datasets_from_function( lambda input_context: make_distributed_dataset(input_context, FILE_PATTERN) ) with strategy.scope(): model = tf.keras.Sequential([ tf.keras.layers.Input(shape=(2049, 2)), # 4096//2+1 个频点，每个频点实部+虚部 tf.keras.layers.Flatten(), tf.keras.layers.Dense(64, activation="relu"), tf.keras.layers.Dense(2, activation="softmax"), ]) model.compile( optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"], ) model.fit(train_ds, epochs=10)这里面最容易被忽略的是input_context.get_per_replica_batch_size(GLOBAL_BATCH_SIZE)。它会把全局 batch size 按照 worker 数量切成每个 worker 自己应处理的副本大小。如果你直接写batch(64)，在多 worker 下每个 worker 都读 64 条，总 batch 就不对了，梯度的 all-reduce 语义也会出问题。另外，文件级分片要求文件数量远大于 worker 数量。如果只有一个大文件，或者文件数量小于 worker 数，会出现某些 worker 分不到数据。这种场景可以改成记录级分片：ds = tf.data.TFRecordDataset(files) ds = ds.shard( input_context.num_input_pipelines, input_context.input_pipeline_id )记录级分片会打散到单条记录，但 IO 性能略差一些，因为每个 worker 都要打开同一批文件去读取自己需要的记录。文件级分片仍然是优先选择。4.3 多 worker 下批次大小、复数和标签对齐的细节先说标签。很多第一次写的人会误以为频谱维度变了，标签也要跟着变。其实标签是对“整条样本”而言的，一个 waveform 对应一个 label，不管 rfft 之后频点数变成多少，label 始终保持不变，注意别在 map 里把 label 的维度不小心压掉就行。再看模型输入。原始波形 8192 点，fft_length=[4096]表示只取前 4096 点参与 FFT，输出 2049 个复数频点，每个频点拆成实部和虚部，最终特征 shape 是(2049, 2)。对应到模型Input(shape=(2049, 2))，这两维千万不要搞反。如果你用的是tf.signal.stft，输出变成(frames, num_freq_bins)，模型输入还要相应调整。最后是复数问题。rfft 输出 complex64，拆成实部和虚部之后变成两个实数张量。手写tf.stack是在管道里处理；如果放在模型层里，同样要拆，否则后面接 Conv2D 或者 Dense 会直接报 dtype 不支持。这个坑我踩了不止一次，尤其从单卡切到多卡时，踩的频率更高。4.4 不训练模型，只想用 TensorFlow 做批量分布式 FFT有一个容易忽视的用法：如果你手头有海量波形需要转成频谱，但暂时不训练模型，完全可以借用 TensorFlow 的分布式能力来跑批量 FFT，没必要再单独起一个大数据工具链。做法很简单：把前面的FFTFeatureLayer包进一个Sequential模型，编译之后直接model.predict(dataset)。prediction 阶段照样走distribute_datasets_from_function做数据分片，每个 worker 在各自的副本上执行 FFT，最后把结果拼回来。这个技巧我在做某次批量频谱转换时用过，省掉了专门的分布式预处理程序。5. 跑起来之后才会注意到的性能细节5.1 num_parallel_calls 不是越大越好tf.data里的map有num_parallel_calls参数，很多人一看并行就调到很大。实测下来并非如此。FFT 本身是计算密集型操作，在 CPU 上跑会占满核心；如果你的num_parallel_calls大于物理核数，线程切换的开销反而会让吞吐下降。我做过一个对比：同一个rfft_to_feature管道，在 8 核机器上把num_parallel_calls从 4 调到 16，每秒处理样本数不升反降，大约降到原来的 70% 左右。我的建议是直接用tf.data.AUTOTUNE，让框架自己根据机器负载去调节。如果非要手动设，先看 CPU 核心数，设置为物理核心数附近，不要盲目拉满。5.2 fft_length、窗函数和频谱泄漏fft_length这个参数很容易被忽视，但它直接影响频谱形态。如果输入长度大于fft_length，tf.signal.rfft会隐式截断，直接丢掉后面多出来的点；如果输入长度小于fft_length，会隐式补零。很多人以为补零只是把点数变多、看起来更平滑，实际上补零并不会增加真实频率分辨率，它只是在连续频谱上做了更多采样，视觉上变平滑而已。更隐蔽的是截断：如果你怀疑高频段丢失了信息，先检查是不是输入长度和fft_length不匹配。靠窗函数减少频谱泄漏也是必修课。实际采集波形往往不是整周期截断，直接 FFT 会导致频谱泄漏，表现为峰附近出现额外旁瓣。在管道里加窗很简单：WINDOW = tf.signal.hann_window(4096, dtype=tf.float32) def rfft_to_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) waveform = waveform[:4096] * WINDOW # 先手动截断，再加窗 spectrum = tf.signal.rfft(waveform, fft_length=[4096]) feature = tf.stack([tf.math.real(spectrum), tf.math.imag(spectrum)], axis=-1) return feature, label注意，窗口长度必须和真正参与 FFT 的点数一致。如果你直接对 8192 点波形加 4096 点的窗，然后期望fft_length=[4096]截断后正好和窗口匹配，这是错的。因为 rfft 是在加窗之后的整个序列上做的，内部先截断还是先加窗，行为和你手写代码时的预期可能不一样。我自己通常是先显式切片到指定长度，再乘窗口，最后才做 rfft，这样每一步都可控。5.3 混合精度与 dtype 问题训练里开了自动混合精度之后，很多层会自动用 float16 计算。但tf.signal.rfft对输入 dtype 有要求，通常只支持 float32 或 float64。如果你在数据管道早期就把波形 cast 成 float16，后面调用 rfft 时要么报错，要么需要框架隐式转换，既慢又损失精度。我的习惯是在管道里统一保证波形是 float32，rfft 输出自然是 complex64。等到频谱拆成实部虚部之后，模型内部是否转 float16 由混合精度策略去决定，不要把 dtype 问题带到 FFT 算子里。5.4 多 worker 下的 IO 与缓存多 worker 训练时，如果每个 epoch 都重新从共享存储读原始波形，IO 压力会非常大。一个折中方案是dataset.cache()：把经过 parse 和初步清洗的数据缓存到本地磁盘，避免每一轮都重复读共享存储。但缓存要慎用，如果数据量大到几个 TB，本地磁盘放不下，缓存就失去意义。文件级分片对 IO 的优化比缓存更直接。因为它避免多个 worker 重复读同一份文件。如果你的文件数量足够多，文件级分片 +interleave+prefetch通常就能把 IO 压到一个可控水平。如果 IO 依然紧张，再考虑离线频谱化。6. 进阶玩法：STFT、二维 FFT 和可学习的频域滤波6.1 STFT：把一维波形变成时间-频率图语音、振动这一类非平稳信号，单次 FFT 会把整个时间段的频率特征混在一起，丢失时间信息。更常见的做法是做 STFT，得到一张时间-频率图，再交给 CNN 去分类。TensorFlow 里tf.signal.stft把分帧、加窗、FFT 合在了一起：def waveform_to_stft_feature(waveform, label): waveform = tf.cast(waveform, tf.float32) spectrogram = tf.signal.stft( waveform, frame_length=256, frame_step=128, fft_length=256, window_fn=tf.signal.hann_window, pad_end=True, ) # 输出 [frames, 129] complex64 spectrogram = tf.abs(spectrogram) spectrogram = tf.math.log1p(spectrogram) # 压缩动态范围 return spectrogram, labelframe_length是窗长，frame_step是帧移，fft_length参与 FFT 的点数。一般建议fft_length等于frame_length，这样窗口函数和 FFT 截断区域完全一致。pad_end=True会在信号尾部补零，保证最后一帧能被处理，避免丢掉尾部信息。STFT 生成的特征比一维频谱大得多。在分布式场景里，如果每个 worker 的 CPU 管道里都跑 STFT，计算量可不小。我的经验是：单样本帧数不多、数据量中等时，放数据管道没问题；如果帧数很多，或者每帧频点很大，可以考虑把 STFT 写进自定义层，让 GPU 去并行。6.2 二维 FFT：图像和雷达数据的维度游戏图像或者雷达距离-多普勒图这类数据，二维 FFT 更常用。tf.signal.rfft2d只作用于最后两维，所以灰度图像[batch, height, width]可以直接用：spectrum = tf.signal.rfft2d(image, fft_length=[height, width]) # 输出 [batch, height, width//2+1] complex64但如果图像是[batch, height, width, channels]，想对每个通道分别做二维 FFT，绝对不能用原始形状直接丢进去，否则rfft2d会对width和channels做变换，结果完全不对。需要先把通道维挪到前面：image_t = tf.transpose(image, [0, 3, 1, 2]) # [batch, channels, height, width] spectrum = tf.signal.rfft2d(image_t) # 对 height/width 做 FFT这个问题本质上还是之前说的“固定作用维度”问题。凡是遇到多维输入，先停下来想清楚：要变换的轴是不是最后两维？如果不是，先 transpose 再操作。6.3 在频谱域做可学习滤波，让 FFT 参与反向传播FFT 是线性算子，TensorFlow 对它也做了梯度支持。利用这一点，可以在模型中间插入一个“频域滤波层”：先把波形 FFT 到频域，乘一个可学习的 mask，再 IFFT 回时域，让模型自己学会在频域上怎么滤波。这在语音增强、去噪这类任务里非常实用。简化实现大致是这样：class LearnableFreqMaskLayer(tf.keras.layers.Layer): def __init__(self, fft_length): super().__init__() self.fft_length = fft_length def build(self, input_shape): self.mask = self.add_weight( name="freq_mask", shape=(self.fft_length // 2 + 1,), initializer="ones", trainable=True, ) def call(self, inputs): spectrum = tf.signal.rfft(inputs, fft_length=[self.fft_length]) spectrum = spectrum * tf.cast(self.mask, spectrum.dtype) return tf.signal.irfft(spectrum, fft_length=[self.fft_length])这个 mask 是模型的一个全局可训练变量，分布式训练时它会参与梯度同步，所以多卡也自然适用。这类层放在模型中间，前面提到的方案 A 就无能为力了，必须用自定义层。这也是我说“先看任务需求再决定 FFT 放哪一层”的原因。最后聊一个实践中的小判断标准。我一般默认把 FFT 放进数据管道，只有当任务本身要求“在频谱上做可学习处理”，也就是需要把频谱的梯度传回原始波形时，才把它放进模型内部。数据管道版本的好处是简单、稳、好排查；模型内部版本的好处是灵活、能端到端优化，但代价是你必须把复数谱拆好，还要面对模型序列化、多 worker 下发等一堆细节。两种方案我都跑过，如果你刚开始搭，建议先从数据管道版本开始，跑通之后再决定要不要往模型里挪。