如何将这段TensorFlow的STFT处理代码转换为PyTorch代码？

原TensorFlow代码的功能是对输入的脑电信号（形状为[samples, num_of_eeg, time_steps]）执行STFT变换，提取频谱图后调整维度顺序，最终输出形状为[samples, time_stamps, freq, num_of_eeg]的张量。以下是对应的PyTorch实现：

import torch

def transforms(x):
    # 输入x形状：[samples, num_of_eeg, time_steps]
    # 执行STFT变换，参数与TensorFlow一一对应
    spectrograms = torch.stft(
        x,
        n_fft=64,
        win_length=32,
        hop_length=4,
        return_complex=True,
        center=False  # 与TensorFlow默认行为保持一致
    )
    # 取复数的模，得到频谱图幅度
    spectrograms = torch.abs(spectrograms)
    # 调整维度顺序：从[samples, num_of_eeg, time_stamps, freq]转为[samples, time_stamps, freq, num_of_eeg]
    spectrograms = spectrograms.permute(0, 2, 3, 1)
    return spectrograms

关键细节说明 #

• 参数对应：PyTorch的torch.stft与TensorFlow的tf.signal.stft参数需对齐：
  • TensorFlow的frame_length → PyTorch的win_length（每帧样本数）
  • TensorFlow的frame_step → PyTorch的hop_length（帧步长）
  • TensorFlow的fft_length → PyTorch的n_fft（FFT运算长度）
• 复数处理：PyTorch 1.7+版本支持return_complex=True，直接返回复数张量，与TensorFlow输出类型对齐，取模操作和tf.abs行为完全一致。
• 维度调整：原TensorFlow用tf.einsum("...ijk->...jki")调整维度，对应PyTorch的permute(0, 2, 3, 1)，将num_of_eeg维度移到最后，同时保留时间步和频率维度的顺序。

内容的提问来源于stack exchange，提问作者longwild laptop