Pytorch自定义快速傅里叶变换层（FFT Layer）

---

前言

在一些网络中，需要实时将时域信号转换为频域信号，快速傅里叶变换（Fast Fourier Transform，FFT）是一种常用方法。
本文主要介绍如何用pytorch实现FFT，并封装成一个自定义层，方便在一些需要时频变换的网络中即插即用。
notebook地址：https://github.com/WKQ9411/FFT_Layer

---

一、构造信号样例

首先构造一个多通道信号，其形状为：(channels,time_samples)，然后转换为tensor格式。代码如下：

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn


# 设置信号参数
num_channels = 5   # 信号通道数
num_samples = 600  # 时间样本点
sampling_rate = 200  # 采样率
low_freq = 0  # 选择想要的频率范围下限
high_freq = 80  # # 选择想要的频率范围上限【注：这里并不是滤波】

# 构造一个多通道信号
t = np.linspace(start=0, stop=(num_samples-1)/sampling_rate, num=num_samples)  # 时间轴
signals = 3*np.sin(2*np.pi*3*t)+2*np.sin(2*np.pi*15*t)+np.sin(2*np.pi*27*t)  # 包含三个分量，幅值分别为3,2,1，频率分别为3,15,27
noise_level = 0.5
noise = np.random.normal(0, noise_level, signals.shape)
signals = signals + noise  # 加入噪声
signals = np.tile(signals, (num_channels, 1))  # 复制成多通道信号
print('signals shape:', signals.shape)  # (channels,time_samples)

# 转换为tensor格式信号
tensor_signals = torch.tensor(signals, dtype=torch.float32)

二、使用torch.fft模块

torch.fft.fft 和 torch.fft.rfft 都是 PyTorch 库中用于进行快速傅里叶变换（FFT）的函数，但它们适用于不同类型的输入数据和有一些关键的区别：

• torch.fft.fft:
  • 处理实数或复数输入，并输出复数结果，包含正频和负频部分
  • 当输入是实数时，输出是对称的，因此含有冗余信息
  • 当输入是复数时，输出是非对称的，完整表示了所有频率的复数幅度和相位信息
• torch.fft.rfft:
  • 专门用于对实数序列进行傅里叶变换，只处理实数输入
  • 输出非负频率的复数结果，从而减少计算复杂度

代码如下：

# 使用PyTorch进行FFT
fft_signals = torch.fft.fft(tensor_signals, dim=1)
rfft_signals = torch.fft.rfft(tensor_signals, dim=1)
print('fft_signals shape:', fft_signals.shape)
print('rfft_signals shape:', rfft_signals.shape)

# 获取频率列表
freqs = torch.fft.fftfreq(num_samples, 1/sampling_rate)  # 得到全部频率
rfreqs = torch.fft.rfftfreq(num_samples, 1/sampling_rate)  # 得到非负频率
print('freqs length:', len(freqs))
print('freqs list:\n', freqs)
print('rfreqs length:', len(rfreqs))
print('rfreqs list:\n', rfreqs)

# 后续步骤使用freqs
# 选择想要的频率范围
mask = (freqs >= low_freq) & (freqs <= high_freq)

# 选择一个通道的信号进行可视化
channel_idx = 0
selected_channel_signal = signals[channel_idx]  # 时域信号
selected_channel_fft = fft_signals[channel_idx][mask]  # 选择频率范围内的频域信号

# 绘制原始信号
plt.figure(figsize=(12, 6))
plt.subplot(2, 1, 1)
plt.plot(t, selected_channel_signal)
plt.title(f'Original Signal of Channel {channel_idx+1}')
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')

# 绘制频域信号
plt.subplot(2, 1, 2)
plt.plot(freqs[mask], 2.0/num_samples * torch.abs(selected_channel_fft))  # 2.0/num_samples因子用于标准化
plt.title(f'FFT of Channel {channel_idx+1} ({low_freq}-{high_freq} Hz)')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Magnitude')

plt.tight_layout()
plt.show()

这里选择第一个通道，绘图结果如下：

三、构造FFT_Layer

为了便于即插即用，封装成一个自定义层，其中low_freq和high_freq用于获取需要的频率范围，方便某些情况下选择期望的频率特征。
代码如下：

# 自定义FFT层
class FFTLayer(nn.Module):
    """
    in: (batch,channel,time)
    out: (bath,channel,freq_len), (freq_len,)
    """

    def __init__(self, sampling_rate, low_freq=None, high_freq=None):
        super().__init__()
        self.low_freq = low_freq
        self.high_freq = high_freq
        self.sampling_rate = sampling_rate

    def forward(self, x):
        num_samples = x.shape[2]
        x = torch.fft.fft(x, dim=2)
        factor = 2.0
        freqs = torch.fft.fftfreq(num_samples, 1 / self.sampling_rate)
        
        if (self.low_freq is None) and (self.high_freq is None):
            fft = torch.abs(x)
        else:
            if (self.low_freq is None) and (self.high_freq is not None):
                mask = freqs <= self.high_freq
            elif (self.low_freq is not None) and (self.high_freq is None):
                mask = freqs >= self.low_freq
            else:
                mask = (freqs >= self.low_freq) & (freqs <= self.high_freq)
            fft = torch.abs(x[..., mask])
            freqs = freqs[mask]
        
        if self.low_freq is None or self.low_freq < 0:  # 如果包含了负频率，标准化因子就不补偿2，而是1
            factor = 1.0
            print(factor)
        return (factor / num_samples) * fft, freqs

使用该自定义层进行FFT，代码如下：

sig = tensor_signals.unsqueeze(0)  # 将(channels,time_samples)转换为(1,channels,time_samples)，其中，1代表batch

fft, freq = FFTLayer(sampling_rate=sampling_rate, low_freq=low_freq, high_freq=high_freq)(sig)
print(fft.shape)
print(freq.shape)

# 绘制原始信号
plt.figure(figsize=(12, 6))
plt.subplot(2, 1, 1)
plt.plot(t, sig[0, channel_idx, :])
plt.title(f'Original Signal of Channel {channel_idx+1}')
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')

# 绘制频域信号
plt.subplot(2, 1, 2)
plt.plot(freq, fft[0, channel_idx, :])
plt.title(f'FFT of Channel {channel_idx+1} ({low_freq}-{high_freq} Hz)')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Magnitude')

plt.tight_layout()
plt.show()

与前文所绘制的图一致。
另外，如果默认不设置low_freq和high_freq，即

fft, freq = FFTLayer(sampling_rate=sampling_rate)(sig)

则绘制全部的正负频率，如下图：