SEGAN 音频降噪模型代码详解及优化建议

本文提供 SEGAN 音频降噪模型的 PyTorch 代码实现，并结合降噪效果分析，提出音频音量减小和高频噪声增加的优化建议，帮助用户更好地理解和改进 SEGAN 模型。

import torch
from dataset import SEGAN_Dataset
from hparams import hparams
from model import Generator, Discriminator
import os
from torch.utils.data import Dataset,DataLoader
import torch.nn as nn
from torch.autograd import Variable

if __name__ == "__main__":
    
   
    
    # 定义device

    device = torch.device("cuda:1")
    
    # 导入参数

    para = hparams()
    
    # 创建数据保存文件夹

    os.makedirs(para.save_path,exist_ok=True)
    
    # 创建生成器

    generator = Generator()
    generator = generator.to(device)
    
    # 创建鉴别器

    discriminator = Discriminator()
    discriminator = discriminator.to(device)
    
    # # 创建G 的优化器

    # m_G_optimizer = torch.optim.RMSprop(m_G.parameters(), lr=para.lr_G)

    
    # # 创建D 的优化器

    # d_optimizer = torch.optim.RMSprop(m_D.parameters(), lr=para.lr_D)

    # optimizers

    
    g_optimizer = torch.optim.RMSprop(generator.parameters(), lr=0.0001)
    d_optimizer = torch.optim.RMSprop(discriminator.parameters(), lr=0.0001)
    
    # 定义数据集

    m_dataset = SEGAN_Dataset(para)
    
    # 获取ref-batch

    ref_batch = m_dataset.ref_batch(para.ref_batch_size)
    ref_batch = Variable(ref_batch)
    ref_batch = ref_batch.to(device)
    
    # 定义dataloader

    m_dataloader = DataLoader(m_dataset,batch_size = para.batch_size,shuffle = True, num_workers = 8)
    loss_d_all =0
    loss_g_all =0
    n_step =0
    for epoch in range(para.n_epoch):
        
        for i_batch, sample_batch in enumerate(m_dataloader):
            batch_clean = sample_batch[0]
            batch_noisy = sample_batch[1]
            batch_clean = Variable(batch_clean)
            batch_noisy = Variable(batch_noisy)
            
            batch_clean = batch_clean.to(device)
            batch_noisy = batch_noisy.to(device)
            
            # print(batch_clean.size())

            # print(batch_noisy.size())

            
            batch_z = nn.init.normal(torch.Tensor(batch_clean.size(0), para.size_z[0], para.size_z[1]))
            batch_z = Variable(batch_z)
            batch_z = batch_z.to(device)
            
            
            discriminator.zero_grad()
            train_batch = Variable(torch.cat([batch_clean,batch_noisy],axis=1))
            # train_batch = torch.cat([batch_clean,batch_noisy],axis=1)

            outputs = discriminator(train_batch, ref_batch)
            clean_loss = torch.mean((outputs - 1.0) ** 2)  # L2 loss - we want them all to be 1

            # clean_loss.backward()


            # TRAIN D to recognize generated audio as noisy

            generated_outputs = generator(batch_noisy, batch_z)
            outputs = discriminator(torch.cat((generated_outputs, batch_noisy), dim=1), ref_batch)
            noisy_loss = torch.mean(outputs ** 2)  # L2 loss - we want them all to be 0

            # noisy_loss.backward()


            d_loss = clean_loss + noisy_loss
            d_loss.backward()
            d_optimizer.step()  # update parameters


            # TRAIN G so that D recognizes G(z) as real

            generator.zero_grad()
            generated_outputs = generator(batch_noisy, batch_z)
            gen_noise_pair = torch.cat((generated_outputs, batch_noisy), dim=1)
            outputs = discriminator(gen_noise_pair, ref_batch)

            g_loss_ = 0.5 * torch.mean((outputs - 1.0) ** 2)
            # L1 loss between generated output and clean sample

            l1_dist = torch.abs(torch.add(generated_outputs, torch.neg(batch_clean)))
            g_cond_loss = 100 * torch.mean(l1_dist)  # conditional loss

            g_loss = g_loss_ + g_cond_loss

            # backprop + optimize

            g_loss.backward()
            g_optimizer.step()
            
            
            
            # # 更新参数 D

            # m_D.zero_grad()

            
            # # 计算真实数据的 d_loss

            # batch_train_real = torch.cat((batch_clean,batch_noisy),axis=1)

  
            
            # real_d_out = m_D(batch_train_real,batch_ref)

            # d_real_loss = torch.mean((real_d_out-1.0)**2)

            # d_real_loss.backward()

            
            # # 计算虚假数据的 d_loss

            # batch_g = m_G(batch_noisy,batch_z)

            # batch_train_fake = torch.cat((batch_g,batch_noisy),axis=1)

            # fake_d_out = m_D(batch_train_fake,batch_ref)

            # d_fake_loss = torch.mean(fake_d_out**2)

            # d_fake_loss.backward()

            
            # d_loss =  d_real_loss + d_fake_loss

            
            # # d_loss.backward()

            # m_D_optimizer.step()

            
            # # 更新参数G

            # m_G.zero_grad()

            
            # batch_g = m_G(batch_noisy,batch_z)

            
            # batch_train_fake = torch.cat((batch_g,batch_noisy),axis=1)

            # fake_d_out = m_D(batch_train_fake,batch_ref)

            # # D 将生成语音认作真实语音

            # g_loss_1 = 0.5*torch.mean((fake_d_out-1.0)**2)

            # # 生成语音和真实语音的L1距离最小

           
            # L1_dis = torch.abs(torc

音频音量减小和高频噪声增加的优化建议

1. 音频音量减小:

原因: 去噪方法可能不够有效，导致一些原本有用的信号也被去除了。
优化建议:
- 调整模型参数: 例如调整 Generator 和 Discriminator 的网络结构、激活函数、损失函数等参数。
- 使用更好的去噪算法: 例如使用更先进的降噪算法，例如基于深度学习的降噪算法或基于信号处理的降噪算法。
- 增加训练数据: 使用更多包含不同类型噪声和不同音频内容的训练数据来训练模型，以提高模型的泛化能力。

2. 高频噪声增加:

原因: 模型过度拟合了训练数据中的高频噪声，不能很好地泛化到其他数据。
优化建议:
- 增加正则化项: 例如在损失函数中添加 L1 或 L2 正则化项，以避免模型过度拟合。
- 使用更多的训练数据: 使用更多包含不同类型噪声和不同音频内容的训练数据来训练模型，以提高模型的泛化能力。
- 添加一些噪声数据来训练模型: 在训练数据中添加一些人工噪声数据，以提高模型对噪声的鲁棒性。

总结

SEGAN 模型在音频降噪方面具有良好的效果，但其性能也取决于模型的训练和参数选择。通过分析降噪效果，并根据具体情况进行调整和尝试，可以进一步提高模型的性能，获得更好的降噪效果。

注意: 以上代码仅供参考，需要根据具体情况进行修改和优化。同时，建议使用更专业的音频处理库来进行音频降噪，例如 Librosa 或 PyAudio，以获得更好的效果。