rewrote demucs model

changed configs default values
refactoring

datasets/valentini.py

@@ -2,10 +2,8 @@ import torch
 from torch.utils.data import Dataset
 from pathlib import Path
 import torchaudio
-import numpy as np
-from torchaudio.transforms import Resample
 
-HIGH_RANDOM_SEED = 1000
+MAX_RANDOM_SEED = 1000
 
 class Valentini(Dataset):
     def __init__(self, dataset_path, val_fraction, transform=None, valid=False, *args, **kwargs):
@@ -34,9 +32,8 @@ class Valentini(Dataset):
         clean_wav, clean_sr = torchaudio.load(self.clean_wavs[idx])
 
         if self.transform:
-            random_seed = 0 if self.valid else torch.randint(HIGH_RANDOM_SEED, (1,))[0]
+            random_seed = 0 if self.valid else torch.randint(MAX_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)

testing/unit_tests.py

@@ -0,0 +1,4 @@
+
+
+def test_model_inference():
+    assert 1 == 1

train.py

@@ -12,21 +12,30 @@ from losses import get_loss
 from datasets import get_datasets
 from testing.metrics import Metrics
 from datasets.minimal import Minimal
+from tqdm import tqdm
 
-
-def train(cfg: DictConfig):
-    device = torch.device(f'cuda:{cfg.gpu}' if torch.cuda.is_available() else 'cpu')
-
+def init_wandb(cfg):
     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))
-    wandb.run.name = cfg['wandb']['run_name']
+               config=omegaconf.OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True),
+               resume=cfg['wandb']['resume'],
+               name=cfg['wandb']['run_name'])
+    if wandb.run.resumed:
+        api = wandb.Api()
+        runs = api.runs(f"{cfg['wandb']['entity']}/{cfg['wandb']['project']}",
+                        order='train_pesq')
+        run = [run for run in runs if run.name == cfg['wandb']['run_name'] and run.state != 'running'][0]
+        artifacts = run.logged_artifacts()
+        best_model = [artifact for artifact in artifacts if artifact.type == 'model'][0]
+
+        best_model.download()
 
+def train(cfg: DictConfig):
+    device = torch.device(f'cuda:{cfg.gpu}' if torch.cuda.is_available() else 'cpu')
+    init_wandb(cfg)
     checkpoint_saver = CheckpointSaver(dirpath=cfg['training']['model_save_path'], run_name=wandb.run.name)
-    metrics = Metrics(rate=cfg['dataloader']['sample_rate'])
+    metrics = Metrics(source_rate=cfg['dataloader']['sample_rate']).to(device)
 
     model = get_model(cfg['model']).to(device)
     optimizer = get_optimizer(model.parameters(), cfg['optimizer'])
@@ -35,14 +44,16 @@ def train(cfg: DictConfig):
     minimal_dataset = Minimal(cfg)
 
     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),
+        'train':  DataLoader(train_dataset, batch_size=cfg['dataloader']['train_batch_size'], shuffle=True,
+                             num_workers=cfg['dataloader']['num_workers']),
+        'val': DataLoader(valid_dataset, batch_size=cfg['dataloader']['valid_batch_size'], shuffle=False,
+                          num_workers=cfg['dataloader']['num_workers']),
         'minimal': DataLoader(minimal_dataset)
     }
 
-    wandb.watch(model, log_freq=100)
-
-    for epoch in range(cfg['training']['num_epochs']):
+    wandb.watch(model, log_freq=cfg['wandb']['log_interval'])
+    epoch = 0
+    while epoch < cfg['training']['num_epochs']:
         for phase in ['train', 'val']:
             if phase == 'train':
                 model.train()
@@ -50,7 +61,8 @@ def train(cfg: DictConfig):
                 model.eval()
 
             running_loss, running_pesq, running_stoi = 0.0, 0.0, 0.0
-            for i, (inputs, labels) in enumerate(dataloaders[phase]):
+            loop = tqdm(dataloaders[phase])
+            for i, (inputs, labels) in enumerate(loop):
                 inputs = inputs.to(device)
                 labels = labels.to(device)
 
@@ -64,11 +76,16 @@ def train(cfg: DictConfig):
                         loss.backward()
                         optimizer.step()
 
-                running_metrics = metrics.calculate(denoised=outputs, clean=labels)
+                running_metrics = metrics(denoised=outputs, clean=labels)
                 running_loss += loss.item() * inputs.size(0)
                 running_pesq += running_metrics['PESQ']
                 running_stoi += running_metrics['STOI']
 
+                loop.set_description(f"Epoch [{epoch}/{cfg['training']['num_epochs']}][{phase}]")
+                loop.set_postfix(loss=running_loss / (i + 1) / inputs.size(0),
+                                 pesq=running_pesq / (i + 1) / inputs.size(0),
+                                 stoi=running_stoi / (i + 1) / inputs.size(0))
+
                 if phase == 'train' and i % cfg['wandb']['log_interval'] == 0:
                     wandb.log({"train_loss": running_loss / (i + 1) / inputs.size(0),
                                "train_pesq": running_pesq / (i + 1) / inputs.size(0),
@@ -84,7 +101,13 @@ def train(cfg: DictConfig):
 
             if phase == 'val':
                 for i, (wav, rate) in enumerate(dataloaders['minimal']):
-                    prediction = model(wav.to(device))
+                    if cfg['dataloader']['normalize']:
+                        std = torch.std(wav)
+                        wav = wav / std
+                        prediction = model(wav.to(device))
+                        prediction = prediction * std
+                    else:
+                        prediction = model(wav.to(device))
                     wandb.log({
                         f"{i}_example": wandb.Audio(
                             prediction.detach().cpu().numpy()[0][0],
@@ -92,6 +115,7 @@ def train(cfg: DictConfig):
 
                 checkpoint_saver(model, epoch, metric_val=eposh_pesq,
                                  optimizer=optimizer, loss=epoch_loss)
+        epoch += 1
 
 
 if __name__ == "__main__":

transforms.py

@@ -8,17 +8,24 @@ from torchvision.transforms import RandomCrop
 class Transform(torch.nn.Module):
     def __init__(
             self,
-            input_sr,
+            input_sample_rate,
             sample_rate,
             max_seconds,
+            normalize,
             *args,
             **kwargs
     ):
         super().__init__()
-        self.resample = Resample(orig_freq=input_sr, new_freq=sample_rate)
+        self.input_sample_rate = input_sample_rate
+        self.sample_rate = sample_rate
+        self.resample = Resample(orig_freq=input_sample_rate, new_freq=sample_rate)
         self.random_crop = RandomCrop((1, int(max_seconds * sample_rate)), pad_if_needed=True)
+        self.normalize = normalize
 
     def forward(self, waveform: torch.Tensor) -> torch.Tensor:
-        resampled = self.resample(waveform)
-        croped = self.random_crop(resampled)
-        return croped
+        if self.input_sample_rate != self.sample_rate:
+            waveform = self.resample(waveform)
+        if self.normalize:
+            waveform = waveform / torch.std(waveform)
+        cropped = self.random_crop(waveform)
+        return cropped

utils.py

@@ -2,6 +2,19 @@ import torchaudio
 import torch
 import matplotlib.pyplot as plt
 from pathlib import Path
+from torch.nn.functional import pad
+
+
+
+def pad_cut_batch_audio(wavs, new_shape):
+    wav_length = wavs.shape[-1]
+    new_length = new_shape[-1]
+
+    if wav_length > new_length:
+        wavs = wavs[:, :, :new_length]
+    elif wav_length < new_length:
+        wavs = pad(wavs, (0, new_length - wav_length))
+    return wavs
 
 
 def collect_valentini_paths(dataset_path):