import torch import torch.nn as nn import numpy as np import librosa import os

定义超参数

num_epochs = 100 learning_rate = 0.001 batch_size = 10

定义模型类

class RNN(nn.Module): def init(self, input_size, hidden_size, num_layers): super(RNN, self).init() self.hidden_size = hidden_size self.num_layers = num_layers self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True) self.fc = nn.Linear(hidden_size, 1)

def forward(self, x):
    h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
    out, _ = self.rnn(x, h0)
    out = self.fc(out)
    return out

加载数据

def load_data(path): data, sr = librosa.load(path, sr=16000) return data, sr

处理数据

def process_data(data, sr): data = librosa.stft(data, n_fft=512, hop_length=256) data = np.abs(data) data = np.log1p(data) data = torch.from_numpy(data).float() return data

定义评估指标

def snr(clean, denoise): clean = clean.flatten() denoise = denoise.flatten() noise = clean - denoise return 10 * np.log10(np.sum(clean ** 2) / np.sum(noise ** 2))

def mos(clean, denoise): clean = clean.flatten() denoise = denoise.flatten() noise = clean - denoise return np.mean(np.exp(-0.1 * np.clip(snr(clean, denoise), -10, 20)))

加载数据集

data_dir = 'data' clean_dir = os.path.join(data_dir, 'clean') noisy_dir = os.path.join(data_dir, 'noisy') clean_files = os.listdir(clean_dir)

定义设备

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

定义模型

model = RNN(257, 128, 2).to(device)

定义损失函数和优化器

criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

训练模型

total_loss = [] for epoch in range(num_epochs): for i in range(0, len(clean_files), batch_size): # 加载数据 clean_batch = [] noisy_batch = [] for j in range(batch_size): if i + j >= len(clean_files): break clean_path = os.path.join(clean_dir, clean_files[i + j]) noisy_path = os.path.join(noisy_dir, clean_files[i + j]) clean, sr = load_data(clean_path) noisy, sr = load_data(noisy_path) clean_data = process_data(clean, sr) noisy_data = process_data(noisy, sr) clean_batch.append(clean_data) noisy_batch.append(noisy_data) clean_batch = torch.stack(clean_batch).to(device) noisy_batch = torch.stack(noisy_batch).to(device)

    # 前向传播
    outputs = model(noisy_batch)
    loss = criterion(outputs, clean_batch)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    total_loss.append(loss.item())

# 打印损失值
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

保存模型

torch.save(model.state_dict(), 'model.ckpt')

加载模型

model = RNN(257, 128, 2).to(device) model.load_state_dict(torch.load('model.ckpt'))

评估模型

snr_total = 0 mos_total = 0 for i in range(len(clean_files)): clean_path = os.path.join(clean_dir, clean_files[i]) noisy_path = os.path.join(noisy_dir, clean_files[i]) clean, sr = load_data(clean_path) noisy, sr = load_data(noisy_path) clean_data = process_data(clean, sr) noisy_data = process_data(noisy, sr) clean_data = clean_data.to(device) noisy_data = noisy_data.to(device)

outputs = model(noisy_data.unsqueeze(0))
outputs = outputs.squeeze()
outputs = outputs.cpu().detach().numpy()
outputs = np.expm1(outputs)
outputs = librosa.istft(outputs, hop_length=256)

snr_total += snr(clean, outputs)
mos_total += mos(clean, outputs)

snr_avg = snr_total / len(clean_files) mos_avg = mos_total / len(clean_files)

print('SNR:', snr_avg) print('MOS:', mos_avg)

需要怎么输入什么样的数据路径内容：需要将干净语音和噪声语音分别存储在两个文件夹中，并且两个文件夹中的文件名需要一一对应。可以参考代码中的 'data' 文件夹和 'clean'、'noisy' 两个子文件夹的存储方式。其中，'clean' 文件夹存储的是干净语音，'noisy' 文件夹存储的是加噪后的语音。在评估模型时，只需要指定数据集所在的路径即可。例如：

data_dir = 'data'
clean_dir = os.path.join(data_dir, 'clean')
noisy_dir = os.path.join(data_dir, 'noisy')

其中，data_dir 是数据集所在的路径，clean_dir 存储的是干净语音的路径，noisy_dir 存储的是加噪后的语音的路径。注意，这里假设两个文件夹中的文件名是一一对应的，即干净语音和噪声语音的文件名相同。