★ FEATURED ARTICLE 做信号处理的人应该都经历过这种事手里一堆波形数据想喂进 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_length256、frame_step128会产生约 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, axis0)那样指定作用轴。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_callstf.data.AUTOTUNE) dataset dataset.batch(64) dataset dataset.prefetch(tf.data.AUTOTUNE)这里有几个顺序性细节是我反复改过才固定的第一shuffle要在map之前。如果先做 FFT 再 shuffle相当于把每个样本的频谱先算出来再打乱FFT 计算量白白浪费还占内存。尤其做 STFT 时先 map 后 shuffle 会让 shuffle 缓冲区里的内存开销爆炸因为每个样本从一维波形变成了二维频谱图。第二rfft的输出是复数如果在 map 里不拆实部虚部后面batch出来的张量是 complex64Keras 模型基本没法直接用。第三拆实部虚部的轴要选对tf.stack([real, imag], axis-1)会把每个频点变成一个长度为 2 的向量最终形状是[num_freq_bins, 2]。如果你希望模型把实虚当成两个分离通道也可以改用axis0或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 参数不用重新生成数据集。缺点是在训练时需要额外计算 FFTIO 带宽也可能成为瓶颈。反过来如果任务已经定型STFT 特征非常大或者共享存储的 IO 实在紧张也可以离线一次性把频谱算好存成 TFRecord训练时直接读频谱。代价是以后每次改参数都要重新生成一遍数据。这个取舍要看项目阶段早期探索阶段一定要存时域模型结构定型、进入大规模重复训练阶段再考虑缓存频谱。4.2 distribute_datasets_from_function 的正确打开方式多 worker 训练时最容易犯的错误是把数据集直接strategy.experimental_distribute_dataset(dataset)然后在每个 worker 上整个数据集都重复载入。更灵活、也更高效的方式是用distribute_datasets_from_function在每个 worker 上按input_context做分片各读各的文件各自做 FFT。我整理过一个完整可跑的骨架核心是文件级分片 管道内 FFTimport 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, shuffleTrue) files files.shard( input_context.num_input_pipelines, input_context.input_pipeline_id ) ds files.interleave( tf.data.TFRecordDataset, num_parallel_callstf.data.AUTOTUNE, deterministicFalse, ) ds ds.map(parse_tfrecord, num_parallel_callstf.data.AUTOTUNE) ds ds.shuffle(10000) ds ds.map(rfft_to_feature, num_parallel_callstf.data.AUTOTUNE) per_replica_batch input_context.get_per_replica_batch_size(GLOBAL_BATCH_SIZE) ds ds.batch(per_replica_batch, drop_remainderTrue) 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//21 个频点每个频点实部虚部 tf.keras.layers.Flatten(), tf.keras.layers.Dense(64, activationrelu), tf.keras.layers.Dense(2, activationsoftmax), ]) model.compile( optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy], ) model.fit(train_ds, epochs10)这里面最容易被忽略的是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_lengthtf.signal.rfft会隐式截断直接丢掉后面多出来的点如果输入长度小于fft_length会隐式补零。很多人以为补零只是把点数变多、看起来更平滑实际上补零并不会增加真实频率分辨率它只是在连续频谱上做了更多采样视觉上变平滑而已。更隐蔽的是截断如果你怀疑高频段丢失了信息先检查是不是输入长度和fft_length不匹配。靠窗函数减少频谱泄漏也是必修课。实际采集波形往往不是整周期截断直接 FFT 会导致频谱泄漏表现为峰附近出现额外旁瓣。在管道里加窗很简单WINDOW tf.signal.hann_window(4096, dtypetf.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 时要么报错要么需要框架隐式转换既慢又损失精度。我的习惯是在管道里统一保证波形是 float32rfft 输出自然是 complex64。等到频谱拆成实部虚部之后模型内部是否转 float16 由混合精度策略去决定不要把 dtype 问题带到 FFT 算子里。5.4 多 worker 下的 IO 与缓存多 worker 训练时如果每个 epoch 都重新从共享存储读原始波形IO 压力会非常大。一个折中方案是dataset.cache()把经过 parse 和初步清洗的数据缓存到本地磁盘避免每一轮都重复读共享存储。但缓存要慎用如果数据量大到几个 TB本地磁盘放不下缓存就失去意义。文件级分片对 IO 的优化比缓存更直接。因为它避免多个 worker 重复读同一份文件。如果你的文件数量足够多文件级分片 interleaveprefetch通常就能把 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_length256, frame_step128, fft_length256, window_fntf.signal.hann_window, pad_endTrue, ) # 输出 [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_endTrue会在信号尾部补零保证最后一帧能被处理避免丢掉尾部信息。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//21] 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( namefreq_mask, shape(self.fft_length // 2 1,), initializerones, trainableTrue, ) 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 下发等一堆细节。两种方案我都跑过如果你刚开始搭建议先从数据管道版本开始跑通之后再决定要不要往模型里挪。
阅读完成 · 觉得有帮助？
