Удаление шума нейросетью в pytoch

import torch
import torch.nn as nn
import torch.optim as optim
from torchaudio.datasets import SPEECHCOMMANDS
from torchaudio.transforms import Spectrogram
import torch.nn.functional as F
 
device = torch.device('cpu')
torch.set_num_threads(7)
 
 
# Определение модели для удаления шума на основе RNN
class DenoisingRNN(nn.Module):
    def __init__(self):
        super(DenoisingRNN, self).__init__()
        self.rnn = nn.RNN(input_size=81, hidden_size=81, num_layers=2, batch_first=True)
 
    def forward(self, x):
        # Изменение формы входа на (batch_size, seq_len, input_size)
        x = x.squeeze(1)
        out, _ = self.rnn(x)
        return out
 
# Загрузка и предобработка тренировочных данных
print("Загрузка и подготовка training data")
train_data = SPEECHCOMMANDS('orig', download=True)
spectrogram = Spectrogram()
specs = []
 
# Нахождение максимальной длины спектрограмм
max_len = 0
for data in train_data:
    waveform = data[0]
    label = data[2]
    spec = spectrogram(waveform)
    specs.append(spec)
    max_len = max(max_len, spec.shape[2])
 
# Дополнение спектрограмм до максимальной длины
padded_specs = []
for spec in specs:
    pad_size = max_len - spec.shape[2]
    padded_spec = F.pad(spec, (0, pad_size))
    padded_specs.append(padded_spec)
print(f"Load training data ok max_len: {max_len}")
 
# Создание DataLoader и формирование пакетов тренировочных данных
noisy_data = [(spec, spec) for spec in padded_specs]
batch_size = 32
train_loader = torch.utils.data.DataLoader(noisy_data, batch_size=batch_size, shuffle=True)
 
# Инициализация модели и функции потерь
model = DenoisingRNN()
model.to(device)
 
criterion = nn.SmoothL1Loss()
optimizer = optim.RMSprop(model.parameters(), lr=0.001)
 
# Цикл обучения
num_epochs = 10
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (input_spec, target_spec) in enumerate(train_loader):
        optimizer.zero_grad()
        
        output_spec = model(input_spec)
        target_spec = target_spec.squeeze(1)  # Приведение целевой переменной к нужному размеру
        loss = criterion(output_spec, target_spec)
        
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
 
        if (i+1) % 20 == 0:
            print(f'Epoch: {epoch+1}, Batch: {i+1}/{len(train_loader)}, Loss: {running_loss/10:.4f}')
            running_loss = 0.0
 
# Сохранение обученной модели и максимальной длины
torch.save({'model_state_dict': model.state_dict(), 'max_len': max_len}, 'denoising_rnn_model.pth')

import torch
import torch.nn as nn
import torch.optim as optim
from torchaudio.transforms import Spectrogram
import torch.nn.functional as F
 
# Определение модели для удаления шума на основе RNN
class DenoisingRNN(nn.Module):
    def __init__(self):
        super(DenoisingRNN, self).__init__()
        self.rnn = nn.RNN(input_size=81, hidden_size=81, num_layers=2, batch_first=True)
 
    def forward(self, x):
        # Изменение формы входа на (batch_size, seq_len, input_size)
        x = x.squeeze(1)
        out, _ = self.rnn(x)
        return out
 
# Загрузка модели и максимальной длины
checkpoint = torch.load('denoising_rnn_model.pth')
model = DenoisingRNN()
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
 
# Загрузка и предобработка аудиофайла
waveform, sample_rate = torchaudio.load('audio_file.wav')
spectrogram = Spectrogram()
spec = spectrogram(waveform)
 
# Применение модели для удаления шума
padded_spec = F.pad(spec, (0, checkpoint['max_len'] - spec.shape[2]))
output_spec = model(padded_spec.unsqueeze(0))
 
# Сохранение результата
denoised_waveform = torchaudio.transforms.InverseSpectrogram()(output_spec)
torchaudio.save('denoised_audio_file.wav', denoised_waveform, sample_rate)