refactored code to work with hydra and wandb

checkpoing_saver.py

@@ -0,0 +1,45 @@
+import os
+import numpy as np
+import logging
+import torch
+import wandb
+
+class CheckpointSaver:
+    def __init__(self, dirpath, decreasing=True, top_n=5):
+        """
+        dirpath: Directory path where to store all model weights
+        decreasing: If decreasing is `True`, then lower metric is better
+        top_n: Total number of models to track based on validation metric value
+        """
+        if not os.path.exists(dirpath): os.makedirs(dirpath)
+        self.dirpath = dirpath
+        self.top_n = top_n
+        self.decreasing = decreasing
+        self.top_model_paths = []
+        self.best_metric_val = np.Inf if decreasing else -np.Inf
+
+    def __call__(self, model, epoch, metric_val):
+        model_path = os.path.join(self.dirpath, 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(
+                f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}, & logging model weights to W&B.")
+            self.best_metric_val = metric_val
+            torch.save(model.state_dict(), model_path)
+            self.log_artifact(f'model-ckpt-epoch-{epoch}.pt', model_path, metric_val)
+            self.top_model_paths.append({'path': model_path, 'score': metric_val})
+            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
+        if len(self.top_model_paths) > self.top_n:
+            self.cleanup()
+
+    def log_artifact(self, filename, model_path, metric_val):
+        artifact = wandb.Artifact(filename, type='model', metadata={'Validation score': metric_val})
+        artifact.add_file(model_path)
+        wandb.run.log_artifact(artifact)
+
+    def cleanup(self):
+        to_remove = self.top_model_paths[self.top_n:]
+        logging.info(f"Removing extra models.. {to_remove}")
+        for o in to_remove:
+            os.remove(o['path'])
+        self.top_model_paths = self.top_model_paths[:self.top_n]

conf/config.yaml

@@ -4,12 +4,29 @@ defaults:
   - loss: mse
   - optimizer: sgd
 
+training:
+  num_epochs: 5
+  model_save_path: /media/public/checkpoints
 
 dataloader:
   max_seconds: 2
   sample_rate: 16000
-  batch_size: 12
+  train_batch_size: 12
+  valid_batch_size: 12
 
+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
 
-augmentations:
-  - random_crop
+
+wandb:
+  project: denoising
+  log_interval: 100
+  api_key: local-e23d01ece807cb31e69b2cf4137e4998e4b9856f
+  host: http://localhost:8080/
+  notes: "Experiment note"
+  tags:
+    - baseline

conf/dataset/valentini.yaml

@@ -1,3 +1,4 @@
-
-name: valentini
-path: /media/public/dataset/denoising/DS_10283_2791/
+valentini:
+ dataset_path: /media/public/datasets/denoising/DS_10283_2791/
+ val_fraction: 0.2
+ sample_rate:  48000

datasets/__init__.py

@@ -0,0 +1,18 @@
+from typing import Tuple
+from torch.utils.data import Dataset
+
+from datasets.valentini import Valentini
+from transforms import Transform
+
+DATASETS_POOL = {
+    'valentini': Valentini
+}
+
+
+
+def get_datasets(cfg) -> Tuple[Dataset, Dataset]:
+    name, dataset_params = list(cfg['dataset'].items())[0]
+    transform = Transform(input_sr=dataset_params['sample_rate'], **cfg['dataloader'])
+    train_dataset = DATASETS_POOL[name](valid=False, transform=transform, **dataset_params)
+    valid_dataset = DATASETS_POOL[name](valid=True, transform=transform, **dataset_params)
+    return train_dataset, valid_dataset

datasets.py → datasets/valentini.py

@@ -1,17 +1,19 @@
 import torch
 from torch.utils.data import Dataset
 from pathlib import Path
-from utils import load_wav
+import torchaudio
+import numpy as np
+from torchaudio.transforms import Resample
 
+HIGH_RANDOM_SEED = 1000
 
 class Valentini(Dataset):
-    def __init__(self, dataset_path='/media/public/dataset/denoising/DS_10283_2791/', transform=None,
-                 valid=False):
+    def __init__(self, dataset_path, val_fraction, transform=None, valid=False, *args, **kwargs):
         clean_path = Path(dataset_path) / 'clean_trainset_56spk_wav'
         noisy_path = Path(dataset_path) / 'noisy_trainset_56spk_wav'
         clean_wavs = list(clean_path.glob("*"))
         noisy_wavs = list(noisy_path.glob("*"))
-        valid_threshold = int(len(clean_wavs) * 0.2)
+        valid_threshold = int(len(clean_wavs) * val_fraction)
         if valid:
             self.clean_wavs = clean_wavs[:valid_threshold]
             self.noisy_wavs = noisy_wavs[:valid_threshold]
@@ -22,16 +24,17 @@ class Valentini(Dataset):
         assert len(self.clean_wavs) == len(self.noisy_wavs)
 
         self.transform = transform
+        self.valid = valid
 
     def __len__(self):
         return len(self.clean_wavs)
 
     def __getitem__(self, idx):
-        noisy_wav = load_wav(self.noisy_wavs[idx])
-        clean_wav = load_wav(self.clean_wavs[idx])
+        noisy_wav, noisy_sr = torchaudio.load(self.noisy_wavs[idx])
+        clean_wav, clean_sr = torchaudio.load(self.clean_wavs[idx])
 
         if self.transform:
-            random_seed = torch.randint(100, (1,))[0]
+            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)
@@ -39,8 +42,5 @@ class Valentini(Dataset):
         return noisy_wav, clean_wav
 
 
-DATASETS_POOL = {
-    'valentini': Valentini
-}
 
 

denoisers/demucs.py

@@ -34,7 +34,7 @@ class Decoder(torch.nn.Module):
         self.glu = torch.nn.GLU(dim=-2)
         self.conv2 = torch.nn.ConvTranspose1d(in_channels=in_channels, out_channels=out_channels,
                                               kernel_size=cfg['conv2']['kernel_size'],
-                                              stride=cfg['conv2']['kernel_size'])
+                                              stride=cfg['conv2']['stride'])
         self.relu = torch.nn.ReLU()
 
     def forward(self, x):

testing/metrics.py

@@ -12,7 +12,7 @@ class Metrics:
         self.snr = SignalNoiseRatio()
         
     def calculate(self, denoised, clean):
-        return {'PESQ': self.nb_pesq(denoised, clean),
-                'STOI': self.stoi(denoised, clean)}
+        return {'PESQ': self.nb_pesq(denoised, clean).item(),
+                'STOI': self.stoi(denoised, clean).item()}
     
 

train.py

@@ -1,94 +1,61 @@
 import os
-from torch.utils.tensorboard import SummaryWriter
 import torch
-from torch.nn import Sequential
 from torch.utils.data import DataLoader
-from datetime import datetime
-from torchvision.transforms import RandomCrop
-from utils import load_wav
-from denoisers.demucs import Demucs
 from pathlib import Path
 from omegaconf import DictConfig
+import wandb
+import torchaudio
 
-from optimizers import OPTIMIZERS_POOL
-from losses import LOSSES
-from datasets import DATASETS_POOL
+from checkpoing_saver import CheckpointSaver
 from denoisers import get_model
 from optimizers import get_optimizer
 from losses import get_loss
-
+from datasets import get_datasets
+from testing.metrics import Metrics
+import omegaconf
 
 os.environ['CUDA_VISIBLE_DEVICES'] = "1"
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-#
-#
-# DATASET_PATH = Path('/media/public/dataset/denoising/DS_10283_2791/')
-# VALID_WAVS = {'hard': 'p257_171.wav',
-#               'medium': 'p232_071.wav',
-#               'easy': 'p232_284.wav'}
-# MAX_SECONDS = 2
-# SAMPLE_RATE = 16000
-#
-# transform = Sequential(RandomCrop((1, int(MAX_SECONDS * SAMPLE_RATE)), pad_if_needed=True))
-#
-# training_loader = DataLoader(Valentini(valid=False, transform=transform), batch_size=12, shuffle=True)
-# validation_loader = DataLoader(Valentini(valid=True, transform=transform), batch_size=12, shuffle=True)
 
+def train(cfg: DictConfig):
+    wandb.login(key=cfg['wandb']['api_key'], host=cfg['wandb']['host'])
+    wandb.init(project=cfg['wandb']['project'],
+               notes=cfg['wandb']['notes'],
+               tags=cfg['wandb']['tags'],
+               config=omegaconf.OmegaConf.to_container(
+        cfg, resolve=True, throw_on_missing=True))
 
+    checkpoint_saver = CheckpointSaver(dirpath=cfg['training']['model_save_path'])
+    metrics = Metrics(rate=cfg['dataloader']['sample_rate'])
 
-def train(cfg: DictConfig):
-    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-    model = get_model(cfg['model'])
+    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)
 
-    writer = SummaryWriter('runs/denoising_trainer_{}'.format(timestamp))
-    epoch_number = 0
-
-    EPOCHS = 5
+    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)
 
-    best_vloss = 1_000_000.
-
-    for tag, wav_path in VALID_WAVS.items():
-        wav = load_wav(DATASET_PATH / 'noisy_testset_wav' / wav_path)
-        writer.add_audio(tag=tag, snd_tensor=wav, sample_rate=SAMPLE_RATE)
-    writer.flush()
-
-    for epoch in range(EPOCHS):
-        print('EPOCH {}:'.format(epoch_number + 1))
+    wandb.watch(model, log_freq=100)
 
+    for epoch in range(cfg['training']['num_epochs']):
         model.train(True)
-
-        running_loss = 0.
-        last_loss = 0.
-
         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()
 
-            running_loss += loss.item()
-            if i % 1000 == 999:
-                last_loss = running_loss / 1000  # loss per batch
-                print('  batch {} loss: {}'.format(i + 1, last_loss))
-                tb_x = epoch_number * len(training_loader) + i + 1
-                writer.add_scalar('Loss/train', last_loss, tb_x)
-                running_loss = 0.
-
-        avg_loss = last_loss
+            if i % cfg['wandb']['log_interval'] == 0:
+                wandb.log({"loss": loss})
 
         model.train(False)
 
-        running_vloss = 0.0
+        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
@@ -96,28 +63,29 @@ def train(cfg: DictConfig):
                 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']
+
 
-            avg_vloss = running_vloss / (i + 1)
-            print('LOSS train {} valid {}'.format(avg_loss, avg_vloss))
+            avg_vloss = running_vloss / len(validation_loader)
+            avg_pesq = running_pesq / len(validation_loader)
+            avg_stoi = running_stoi / len(validation_loader)
 
-            writer.add_scalars('Training vs. Validation Loss',
-                               {'Training': avg_loss, 'Validation': avg_vloss},
-                               epoch_number + 1)
-            for tag, wav_path in VALID_WAVS.items():
-                wav = load_wav(DATASET_PATH / 'noisy_testset_wav' / wav_path)
+            wandb.log({"valid_loss": avg_vloss,
+                       "valid_pesq": avg_pesq,
+                       "valid_stoi": avg_stoi})
+
+            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)
-                writer.add_audio(tag=f"Model predicted {tag} on epoch {epoch}",
-                                 snd_tensor=prediction,
-                                 sample_rate=SAMPLE_RATE)
-            writer.flush()
-
-            if avg_vloss < best_vloss:
-                best_vloss = avg_vloss
-                model_path = 'checkpoints/model_{}_{}'.format(timestamp, epoch_number)
-                torch.save(model.state_dict(), model_path)
+                wandb.log({
+                    f"{tag}_epoch_{epoch}": wandb.Audio(
+                        prediction.cpu()[0][0],
+                        sample_rate=rate)})
 
-            epoch_number += 1
+            checkpoint_saver(model, epoch, metric_val=avg_pesq)
 
 
 if __name__ == '__main__':

transforms.py

@@ -0,0 +1,22 @@
+import torch
+
+from torchaudio.transforms import Resample
+from torchvision.transforms import RandomCrop
+
+class Transform(torch.nn.Module):
+    def __init__(
+            self,
+            input_sr,
+            sample_rate,
+            max_seconds,
+            *args,
+            **kwargs
+    ):
+        super().__init__()
+        self.resample = Resample(orig_freq=input_sr, new_freq=sample_rate)
+        self.random_crop = RandomCrop((1, int(max_seconds * sample_rate)), pad_if_needed=True)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        resampled = self.resample(waveform)
+        croped = self.random_crop(resampled)
+        return croped

utils.py

@@ -14,11 +14,6 @@ def collect_valentini_paths(dataset_path):
     return clean_wavs, noisy_wavs
 
 
-def load_wav(path):
-    wav, org_sr = torchaudio.load(path)
-    wav = torchaudio.functional.resample(wav, orig_freq=org_sr, new_freq=16000)
-    return wav
-
 
 
 def plot_spectrogram(stft, title="Spectrogram", xlim=None):