import torch class Network(torch.nn.Module): def __init__(self, n_fft=1024, n_hop=160, n_hidden=1024): super().__init__() self.linear1 = torch.nn.LSTM(n_fft//2+1, n_hidden//2, num_layers=2, batch_first=True, bidirectional=True) self.linear2 = torch.nn.Linear(n_hidden, n_fft//2+1) self.n_fft = n_fft self.n_hop = n_hop # self.window = self.register_buffer('window', torch.hann_window(n_fft)) def forward(self, noisy): # 傅里叶变换 noisy_spec = torch.stft(noisy, self.n_fft, self.n_hop, window=self.window, return_complex=True) # [B,C,T] noisy_amplitude = torch.abs(noisy_spec).transpose(2,1) # [B,T,C=513] # 神经网络计算 hidden = self.linear1(noisy_amplitude)[0] # [B,T,n_hidden=1024] mask = self.linear2(hidden).sigmoid() # [B,T,C=513] (0,1) # 降噪结果 denoise_spec = noisy_spec * mask.transpose(2,1) denoise_wav = torch.istft(denoise_spec, self.n_fft, self.n_hop, window=self.window) # 返回降噪后波形和降噪后谱图 return denoise_wav, denoise_spec if __name__ == "__main__": net = Network() clean = torch.randn(4, 16000) noise = torch.randn(4, 16000) denoise_wav, denoise_spec = net(clean + noise) print(denoise_spec.shape, denoise_wav.shape)

这段代码实现了一个基于LSTM的音频降噪网络，其核心思路是使用LSTM对音频信号的幅度谱进行处理，从而得到一个掩膜（mask），该掩膜可以在频域上降噪信号。具体来说，该网络的输入是含有噪声的干净音频信号，经过傅里叶变换得到幅度谱，之后通过LSTM计算得到一个掩膜，最后将该掩膜应用于幅度谱上，得到降噪后的幅度谱，再通过逆傅里叶变换得到降噪后的音频信号。

具体来说，该网络包括以下几个部分：

- `__init__`函数：定义了一个两层的LSTM网络和一个线性层（用于将LSTM输出的隐藏状态转换为掩膜），并初始化了窗函数等参数。
- `forward`函数：实现了整个网络的前向传播过程。首先对输入信号进行傅里叶变换，得到幅度谱，之后将幅度谱通过LSTM处理得到一个掩膜，再将该掩膜应用于幅度谱上，得到降噪后的幅度谱，最后通过逆傅里叶变换得到降噪后的音频信号。
- `if __name__ == "__main__"`：对网络的功能进行测试，输入为一个随机产生的4秒钟的干净音频信号和一个随机产生的4秒钟的噪声信号，输出为降噪后的幅度谱和音频信号。

需要注意的是，由于该网络使用了LSTM，其计算成本较高，需要在计算资源允许的情况下使用。此外，该网络的性能也取决于数据集的质量和训练过程的参数设置等因素。