refactored train loop

added Multi_STFT loss from paper
added minimal dataset

README.md

@@ -1,3 +1,11 @@
+# MVP
+Сервисом является web interface, в котором пользователь 
+сможет записать своей голос в шумных условиях и получить на выход аудиозапись без шума.
+Для обработки шумных аудио файлов есть доступ к  API на питоне.
+
+Web interface реализован на gradio. Сама работа пишется в контексте фрейморка pytorch.
+В качестве системы контроля экспериментов выбран wandb. Для управления конфигами - hydra.
+Архитектура модели базируется на работе "Real Time Speech Enhancement in the Waveform Domain" от facebook.
 
 
 
@@ -7,4 +15,3 @@
 | ideal denoising |     1.9709     |     0.9211     |
 |    baseline     |     1.7433     |     0.8844     |
 
-

checkpoing_saver.py

@@ -5,7 +5,7 @@ import torch
 import wandb
 
 class CheckpointSaver:
-    def __init__(self, dirpath, decreasing=True, top_n=5):
+    def __init__(self, dirpath, run_name='', decreasing=True, top_n=5):
         """
         dirpath: Directory path where to store all model weights
         decreasing: If decreasing is `True`, then lower metric is better
@@ -17,9 +17,10 @@ class CheckpointSaver:
         self.decreasing = decreasing
         self.top_model_paths = []
         self.best_metric_val = np.Inf if decreasing else -np.Inf
+        self.run_name = run_name
 
     def __call__(self, model, epoch, metric_val):
-        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt')
+        model_path = os.path.join(self.dirpath, self.run_name, model.__class__.__name__ + f'_epoch{epoch}.pt')
         save = metric_val < self.best_metric_val if self.decreasing else metric_val > self.best_metric_val
         if save:
             logging.info(

conf/config.yaml

@@ -16,10 +16,7 @@ dataloader:
 
 validation:
   path: /media/public/datasets/denoising/DS_10283_2791/noisy_testset_wav
-  wavs:
-    easy: p232_284.wav
-    medium: p232_071.wav
-    hard : p257_171.wav
+  sample_rate: 48000
 
 
 wandb:

conf/loss/L1_Multi_STFT.yaml

@@ -0,0 +1 @@
+name: L1_Multi_STFT

datasets/minimal.py

@@ -0,0 +1,24 @@
+import torch
+from torch.utils.data import Dataset
+from pathlib import Path
+import torchaudio
+import numpy as np
+from torchaudio.transforms import Resample
+
+
+class Minimal(Dataset):
+    def __init__(self, cfg):
+        self.wavs = ['p232_284.wav', 'p232_071.wav', 'p257_171.wav']
+        self.dataset_path = cfg['validation']['path']
+        self.target_rate = cfg['dataloader']['sample_rate']
+        self.resampler = Resample(orig_freq=cfg['validation']['sample_rate'],
+                                  new_freq=cfg['dataloader']['sample_rate'])
+
+    def __len__(self):
+        return len(self.wavs)
+
+    def __getitem__(self, idx):
+        wav, rate = torchaudio.load(self.wavs[idx])
+        wav = self.resampler(wav)
+        wav = torch.reshape(wav, (1, 1, -1))
+        return wav, self.target_rate

datasets/valentini.py

@@ -36,6 +36,7 @@ class Valentini(Dataset):
         if self.transform:
             random_seed = 0 if self.valid else torch.randint(HIGH_RANDOM_SEED, (1,))[0]
             torch.manual_seed(random_seed)
+
             noisy_wav = self.transform(noisy_wav)
             torch.manual_seed(random_seed)
             clean_wav = self.transform(clean_wav)

losses.py

@@ -1,7 +1,180 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Original copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+
 import torch
+import torch.nn.functional as F
+
+"""STFT-based Loss modules."""
+
+
+def stft(x, fft_size, hop_size, win_length, window):
+    """Perform STFT and convert to magnitude spectrogram.
+    Args:
+        x (Tensor): Input signal tensor (B, T).
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length.
+        window (str): Window function type.
+    Returns:
+        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
+    """
+    x_stft = torch.stft(x, fft_size, hop_size, win_length, window)
+    real = x_stft[..., 0]
+    imag = x_stft[..., 1]
+
+    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
+    return torch.sqrt(torch.clamp(real ** 2 + imag ** 2, min=1e-7)).transpose(2, 1)
+
+
+class SpectralConvergengeLoss(torch.nn.Module):
+    """Spectral convergence loss module."""
+
+    def __init__(self):
+        """Initilize spectral convergence loss module."""
+        super(SpectralConvergengeLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """Calculate forward propagation.
+        Args:
+            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
+        Returns:
+            Tensor: Spectral convergence loss value.
+        """
+        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
+
+
+class LogSTFTMagnitudeLoss(torch.nn.Module):
+    """Log STFT magnitude loss module."""
+
+    def __init__(self):
+        """Initilize los STFT magnitude loss module."""
+        super(LogSTFTMagnitudeLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """Calculate forward propagation.
+        Args:
+            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
+        Returns:
+            Tensor: Log STFT magnitude loss value.
+        """
+        return F.l1_loss(torch.log(y_mag), torch.log(x_mag))
+
+
+class STFTLoss(torch.nn.Module):
+    """STFT loss module."""
+
+    def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window"):
+        """Initialize STFT loss module."""
+        super(STFTLoss, self).__init__()
+        self.fft_size = fft_size
+        self.shift_size = shift_size
+        self.win_length = win_length
+        self.register_buffer("window", getattr(torch, window)(win_length))
+        self.spectral_convergenge_loss = SpectralConvergengeLoss()
+        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss()
+
+    def forward(self, x, y):
+        """Calculate forward propagation.
+        Args:
+            x (Tensor): Predicted signal (B, T).
+            y (Tensor): Groundtruth signal (B, T).
+        Returns:
+            Tensor: Spectral convergence loss value.
+            Tensor: Log STFT magnitude loss value.
+        """
+        x_mag = stft(x, self.fft_size, self.shift_size, self.win_length, self.window)
+        y_mag = stft(y, self.fft_size, self.shift_size, self.win_length, self.window)
+        sc_loss = self.spectral_convergenge_loss(x_mag, y_mag)
+        mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)
+
+        return sc_loss, mag_loss
+
+
+class MultiResolutionSTFTLoss(torch.nn.Module):
+    """Multi resolution STFT loss module."""
+
+    def __init__(self,
+                 fft_sizes=[1024, 2048, 512],
+                 hop_sizes=[120, 240, 50],
+                 win_lengths=[600, 1200, 240],
+                 window="hann_window", factor_sc=0.1, factor_mag=0.1):
+        """Initialize Multi resolution STFT loss module.
+        Args:
+            fft_sizes (list): List of FFT sizes.
+            hop_sizes (list): List of hop sizes.
+            win_lengths (list): List of window lengths.
+            window (str): Window function type.
+            factor (float): a balancing factor across different losses.
+        """
+        super(MultiResolutionSTFTLoss, self).__init__()
+        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
+        self.stft_losses = torch.nn.ModuleList()
+        for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths):
+            self.stft_losses += [STFTLoss(fs, ss, wl, window)]
+        self.factor_sc = factor_sc
+        self.factor_mag = factor_mag
+
+    def forward(self, x, y):
+        """Calculate forward propagation.
+        Args:
+            x (Tensor): Predicted signal (B, T).
+            y (Tensor): Groundtruth signal (B, T).
+        Returns:
+            Tensor: Multi resolution spectral convergence loss value.
+            Tensor: Multi resolution log STFT magnitude loss value.
+        """
+        sc_loss = 0.0
+        mag_loss = 0.0
+        for f in self.stft_losses:
+            sc_l, mag_l = f(x, y)
+            sc_loss += sc_l
+            mag_loss += mag_l
+        sc_loss /= len(self.stft_losses)
+        mag_loss /= len(self.stft_losses)
+
+        return self.factor_sc*sc_loss, self.factor_mag*mag_loss
+
+
+
+
+class L1_Multi_STFT(torch.nn.Module):
+    """STFT loss module."""
+
+    def __init__(self):
+        """Initialize STFT loss module."""
+        super(L1_Multi_STFT, self).__init__()
+        self.multi_STFT_loss = MultiResolutionSTFTLoss()
+        self.l1_loss =  torch.nn.L1Loss()
+
+    def forward(self, x, y):
+        """Calculate forward propagation.
+        Args:
+            x (Tensor): Predicted signal (B, T).
+            y (Tensor): Groundtruth signal (B, T).
+        Returns:
+            Tensor: Spectral convergence loss value.
+            Tensor: Log STFT magnitude loss value.
+        """
+        sc_loss, mag_loss = self.multi_STFT_loss(x, y)
+        l1_loss = self.l1_loss(x, y)
+        return sc_loss + mag_loss + l1_loss
+
 
 LOSSES = {
-    'mse': torch.nn.MSELoss()
+    'mse': torch.nn.MSELoss(),
+    'L1': torch.nn.L1Loss(),
+    'Multi_STFT': MultiResolutionSTFTLoss,
+    'L1_Multi_STFT': L1_Multi_STFT
 }
 
 

main.py

@@ -5,7 +5,6 @@ from train import train
 
 @hydra.main(version_base=None, config_path="conf", config_name="config")
 def main(cfg: DictConfig):
-    print(OmegaConf.to_yaml(cfg))
     train(cfg)
 
 

notebooks/EDA.ipynb

@@ -2,44 +2,25 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "f800718e-c29f-44d8-bf41-e02d50d0f730",
    "metadata": {
     "ExecuteTime": {
      "start_time": "2023-04-29T13:11:15.198687Z",
      "end_time": "2023-04-29T13:11:15.245584Z"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/maksim/.local/lib/python3.10/site-packages/torchaudio/compliance/kaldi.py:22: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)\n",
-      "  EPSILON = torch.tensor(torch.finfo(torch.float).eps)\n"
-     ]
     },
-    {
-     "ename": "ModuleNotFoundError",
-     "evalue": "No module named 'matplotlib'",
-     "output_type": "error",
-     "traceback": [
-      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
-      "\u001B[0;31mModuleNotFoundError\u001B[0m                       Traceback (most recent call last)",
-      "Cell \u001B[0;32mIn[1], line 3\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mpathlib\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m Path\n\u001B[0;32m----> 3\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mdatasets\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m Valentini\n\u001B[1;32m      5\u001B[0m dataset \u001B[38;5;241m=\u001B[39m Valentini()\n",
-      "File \u001B[0;32m~/PycharmProjects/denoising/datasets.py:4\u001B[0m\n\u001B[1;32m      2\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mutils\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mdata\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m Dataset\n\u001B[1;32m      3\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mpathlib\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m Path\n\u001B[0;32m----> 4\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mutils\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m load_wav\n\u001B[1;32m      7\u001B[0m \u001B[38;5;28;01mclass\u001B[39;00m \u001B[38;5;21;01mValentini\u001B[39;00m(Dataset):\n\u001B[1;32m      8\u001B[0m     \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m__init__\u001B[39m(\u001B[38;5;28mself\u001B[39m, dataset_path\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m/media/public/datasets/denoising/DS_10283_2791/\u001B[39m\u001B[38;5;124m'\u001B[39m, transform\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mNone\u001B[39;00m,\n\u001B[1;32m      9\u001B[0m                  valid\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m):\n",
-      "File \u001B[0;32m~/PycharmProjects/denoising/utils.py:3\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorchaudio\u001B[39;00m\n\u001B[1;32m      2\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\n\u001B[0;32m----> 3\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mmatplotlib\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mpyplot\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m \u001B[38;5;21;01mplt\u001B[39;00m\n\u001B[1;32m      4\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mpathlib\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m Path\n\u001B[1;32m      7\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mcollect_valentini_paths\u001B[39m(dataset_path):\n",
-      "\u001B[0;31mModuleNotFoundError\u001B[0m: No module named 'matplotlib'"
-     ]
+    "pycharm": {
+     "is_executing": true
     }
-   ],
+   },
+   "outputs": [],
    "source": [
     "\n",
     "from pathlib import Path\n",
     "\n",
     "from datasets import Valentini\n",
     "\n",
-    "dataset = Valentini()"
+    "dataset = Valentini('/media/public/datasets/denoising/DS_10283_2791/', valid=False)"
    ]
   },
   {

train.py

@@ -1,10 +1,9 @@
 import os
 import torch
 from torch.utils.data import DataLoader
-from pathlib import Path
+import omegaconf
 from omegaconf import DictConfig
 import wandb
-import torchaudio
 
 from checkpoing_saver import CheckpointSaver
 from denoisers import get_model
@@ -12,7 +11,7 @@ from optimizers import get_optimizer
 from losses import get_loss
 from datasets import get_datasets
 from testing.metrics import Metrics
-import omegaconf
+from datasets.minimal import Minimal
 
 os.environ['CUDA_VISIBLE_DEVICES'] = "1"
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -26,67 +25,68 @@ def train(cfg: DictConfig):
                config=omegaconf.OmegaConf.to_container(
         cfg, resolve=True, throw_on_missing=True))
 
-    checkpoint_saver = CheckpointSaver(dirpath=cfg['training']['model_save_path'])
+    checkpoint_saver = CheckpointSaver(dirpath=cfg['training']['model_save_path'], run_name=wandb.run.name)
     metrics = Metrics(rate=cfg['dataloader']['sample_rate'])
 
     model = get_model(cfg['model']).to(device)
     optimizer = get_optimizer(model.parameters(), cfg['optimizer'])
     loss_fn = get_loss(cfg['loss'])
     train_dataset, valid_dataset = get_datasets(cfg)
+    minimal_dataset = Minimal(cfg)
 
-    training_loader = DataLoader(train_dataset, batch_size=cfg['dataloader']['train_batch_size'], shuffle=True)
-    validation_loader = DataLoader(valid_dataset, batch_size=cfg['dataloader']['valid_batch_size'], shuffle=True)
+    dataloaders = {
+        'train':  DataLoader(train_dataset, batch_size=cfg['dataloader']['train_batch_size'], shuffle=True),
+        'val': DataLoader(valid_dataset, batch_size=cfg['dataloader']['valid_batch_size'], shuffle=True),
+        'minimal': DataLoader(minimal_dataset)
+    }
 
     wandb.watch(model, log_freq=100)
 
     for epoch in range(cfg['training']['num_epochs']):
-        model.train(True)
-        for i, data in enumerate(training_loader):
-            inputs, labels = data
-            inputs, labels = inputs.to(device), labels.to(device)
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = loss_fn(outputs, labels)
-            loss.backward()
-            optimizer.step()
-
-            if i % cfg['wandb']['log_interval'] == 0:
-                wandb.log({"loss": loss})
-
-        model.train(False)
-
-        running_vloss, running_pesq, running_stoi = 0.0, 0.0, 0.0
-        with torch.no_grad():
-            for i, vdata in enumerate(validation_loader):
-                vinputs, vlabels = vdata
-                vinputs, vlabels = vinputs.to(device), vlabels.to(device)
-                voutputs = model(vinputs)
-                vloss = loss_fn(voutputs, vlabels)
-                running_vloss += vloss
-                running_metrics = metrics.calculate(denoised=voutputs, clean=vlabels)
-                running_pesq += running_metrics['PESQ']
-                running_stoi += running_metrics['STOI']
+        for phase in ['train', 'val']:
+            if phase == 'train':
+                model.train()
+            else:
+                model.eval()
 
+            running_loss, running_pesq, running_stoi = 0.0, 0.0, 0.0
+            for i, (inputs, labels) in enumerate(dataloaders[phase]):
+                inputs = inputs.to(device)
+                labels = labels.to(device)
 
-            avg_vloss = running_vloss / len(validation_loader)
-            avg_pesq = running_pesq / len(validation_loader)
-            avg_stoi = running_stoi / len(validation_loader)
+                optimizer.zero_grad()
 
-            wandb.log({"valid_loss": avg_vloss,
-                       "valid_pesq": avg_pesq,
-                       "valid_stoi": avg_stoi})
+                with torch.set_grad_enabled(phase == 'train'):
+                    outputs = model(inputs)
+                    loss = loss_fn(outputs, labels)
 
-            for tag, wav_path in cfg['validation']['wavs'].items():
-                wav, rate = torchaudio.load(Path(cfg['validation']['path']) / wav_path)
-                wav = torch.reshape(wav, (1, 1, -1)).to(device)
-                prediction = model(wav)
-                wandb.log({
-                    f"{tag}_epoch_{epoch}": wandb.Audio(
-                        prediction.cpu()[0][0],
-                        sample_rate=rate)})
-
-            checkpoint_saver(model, epoch, metric_val=avg_pesq)
+                    if phase == 'train':
+                        loss.backward()
+                        optimizer.step()
 
+                running_metrics = metrics.calculate(denoised=outputs, clean=labels)
+                running_loss += loss.item() * inputs.size(0)
+                running_pesq += running_metrics['PESQ']
+                running_stoi += running_metrics['STOI']
 
-if __name__ == '__main__':
-    train()
+                if phase == 'train' and i % cfg['wandb']['log_interval'] == 0:
+                    wandb.log({"train_loss": running_loss / (i + 1),
+                               "train_pesq": running_pesq / (i + 1),
+                               "train_stoi": running_stoi / (i + 1)})
+            epoch_loss = running_loss / len(dataloaders[phase])
+            eposh_pesq = running_pesq / len(dataloaders[phase])
+            eposh_stoi = running_stoi / len(dataloaders[phase])
+
+            wandb.log({f"{phase}_loss": epoch_loss,
+                       f"{phase}_pesq": eposh_pesq,
+                       f"{phase}_stoi": eposh_stoi})
+
+            if phase == 'val':
+                for i, (wav, rate) in enumerate(dataloaders['minimal']):
+                    prediction = model(wav)
+                    wandb.log({
+                        f"{i}_example": wandb.Audio(
+                            prediction.cpu()[0][0],
+                            sample_rate=rate)})
+
+                checkpoint_saver(model, epoch, metric_val=eposh_pesq)

transforms.py

@@ -3,6 +3,8 @@ import torch
 from torchaudio.transforms import Resample
 from torchvision.transforms import RandomCrop
 
+
+
 class Transform(torch.nn.Module):
     def __init__(
             self,