Add vocalset and FAD

config.yaml

@@ -19,13 +19,19 @@ callbacks:
     filename: '{epoch:02d}-{valid_loss:.3f}'
 
 datamodule:
-  _target_: remfx.datasets.Datamodule
-  dataset:
-    _target_: remfx.datasets.GuitarSet
+  _target_: remfx.datasets.VocalSetDatamodule
+  train_dataset:
+    _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size_in_sec: 6
-  val_split: 0.2
+    mode: "train"
+  val_dataset:
+    _target_: remfx.datasets.VocalSet
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size_in_sec: 6
+    mode: "val"
   batch_size: 16
   num_workers: 8
   pin_memory: True

config_guitarset.yaml

@@ -0,0 +1,52 @@
+defaults:
+  - _self_
+  - exp: null
+seed: 12345
+train: True
+sample_rate: 48000
+logs_dir: "./logs"
+log_every_n_steps: 1000
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: "valid_loss"   # name of the logged metric which determines when model is improving
+    save_top_k: 1           # save k best models (determined by above metric)
+    save_last: True         # additionaly always save model from last epoch
+    mode: "min"             # can be "max" or "min"
+    verbose: False
+    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+    filename: '{epoch:02d}-{valid_loss:.3f}'
+
+datamodule:
+  _target_: remfx.datasets.Datamodule
+  dataset:
+    _target_: remfx.datasets.GuitarSet
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size_in_sec: 6
+  val_split: 0.2
+  batch_size: 16
+  num_workers: 8
+  pin_memory: True
+  persistent_workers: True
+
+logger:
+  _target_: pytorch_lightning.loggers.WandbLogger
+  project: ${oc.env:WANDB_PROJECT}
+  entity: ${oc.env:WANDB_ENTITY}
+  # offline: False  # set True to store all logs only locally
+  job_type: "train"
+  group: ""
+  save_dir: "."
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  precision: 32 # Precision used for tensors, default `32`
+  min_epochs: 0
+  max_epochs: -1
+  enable_model_summary: False
+  log_every_n_steps: 1 # Logs metrics every N batches
+  accumulate_grad_batches: 1
+  accelerator: null
+  devices: 1

exp/umx.yaml

@@ -15,7 +15,14 @@ model:
     alpha: 0.3
     sample_rate: ${sample_rate}
 datamodule:
-  dataset:
+  train_dataset:
+    effect_types:
+        Distortion:
+          _target_: remfx.effects.RandomPedalboardDistortion
+          sample_rate: ${sample_rate}
+          min_drive_db: -10
+          max_drive_db: 50
+  val_dataset:
     effect_types:
         Distortion:
           _target_: remfx.effects.RandomPedalboardDistortion

remfx/datasets.py

@@ -20,6 +20,7 @@ from pedalboard import (
 
 # https://zenodo.org/record/7044411/ -> GuitarFX
 # https://zenodo.org/record/3371780  -> GuitarSet
+# https://zenodo.org/record/1193957 -> VocalSet
 
 deterministic_effects = {
     "Distortion": Pedalboard([Distortion()]),
@@ -173,6 +174,74 @@ class GuitarSet(Dataset):
         return (normalized_input, normalized_target, effect_name)
 
 
+class VocalSet(Dataset):
+    def __init__(
+        self,
+        root: str,
+        sample_rate: int,
+        chunk_size_in_sec: int = 3,
+        effect_types: List[torch.nn.Module] = None,
+        mode: str = "train",
+    ):
+        super().__init__()
+        self.chunks = []
+        self.song_idx = []
+        self.root = Path(root)
+        self.chunk_size_in_sec = chunk_size_in_sec
+        self.sample_rate = sample_rate
+        self.mode = mode
+
+        mode_path = self.root / self.mode
+        self.files = sorted(list(mode_path.glob("./**/*.wav")))
+        for i, audio_file in enumerate(self.files):
+            chunk_starts, orig_sr = create_sequential_chunks(
+                audio_file, self.chunk_size_in_sec
+            )
+            self.chunks += chunk_starts
+            self.song_idx += [i] * len(chunk_starts)
+        print(f"Found {len(self.files)} files .\n" f"Total chunks: {len(self.chunks)}")
+        self.resampler = T.Resample(orig_sr, sample_rate)
+        self.effect_types = effect_types
+        self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
+
+    def __len__(self):
+        return len(self.chunks)
+
+    def __getitem__(self, idx):
+        # Load and effect audio
+        song_idx = self.song_idx[idx]
+        x, sr = torchaudio.load(self.files[song_idx])
+        chunk_start = self.chunks[idx]
+        chunk_size_in_samples = self.chunk_size_in_sec * sr
+        x = x[:, chunk_start : chunk_start + chunk_size_in_samples]
+        resampled_x = self.resampler(x)
+        # Reset chunk size to be new sample rate
+        chunk_size_in_samples = self.chunk_size_in_sec * self.sample_rate
+        # Pad to chunk_size if needed
+        if resampled_x.shape[-1] < chunk_size_in_samples:
+            resampled_x = F.pad(
+                resampled_x, (0, chunk_size_in_samples - resampled_x.shape[1])
+            )
+
+        # Add random effect if train
+        if self.mode == "train":
+            random_effect_idx = torch.rand(1).item() * len(self.effect_types.keys())
+            effect_name = list(self.effect_types.keys())[int(random_effect_idx)]
+            effect = self.effect_types[effect_name]
+            effected_input = effect(resampled_x)
+        else:
+            # deterministic static effect for eval
+            effect_idx = idx % len(self.effect_types.keys())
+            effect_name = list(self.effect_types.keys())[effect_idx]
+            effect = deterministic_effects[effect_name]
+            effected_input = torch.from_numpy(
+                effect(resampled_x.numpy(), self.sample_rate)
+            )
+        normalized_input = self.normalize(effected_input)
+        normalized_target = self.normalize(resampled_x)
+        return (normalized_input, normalized_target, effect_name)
+
+
 def create_random_chunks(
     audio_file: str, chunk_size: int, num_chunks: int
 ) -> Tuple[List[Tuple[int, int]], int]:
@@ -249,3 +318,43 @@ class Datamodule(pl.LightningDataModule):
             pin_memory=self.pin_memory,
             shuffle=False,
         )
+
+
+class VocalSetDatamodule(pl.LightningDataModule):
+    def __init__(
+        self,
+        train_dataset,
+        val_dataset,
+        *,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool = False,
+        **kwargs: int,
+    ) -> None:
+        super().__init__()
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+
+    def setup(self, stage: Any = None) -> None:
+        pass
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+            dataset=self.val_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+        )

remfx/models.py

@@ -7,11 +7,43 @@ from audio_diffusion_pytorch import DiffusionModel
 from auraloss.time import SISDRLoss
 from auraloss.freq import MultiResolutionSTFTLoss, STFTLoss
 from torch.nn import L1Loss
+from frechet_audio_distance import FrechetAudioDistance
+import numpy as np
 
 from umx.openunmix.model import OpenUnmix, Separator
 from torchaudio.models import HDemucs
 
 
+class FADLoss(torch.nn.Module):
+    def __init__(self, sample_rate: float):
+        super().__init__()
+        self.fad = FrechetAudioDistance(
+            use_pca=False, use_activation=False, verbose=False
+        )
+        self.sr = sample_rate
+
+    def forward(self, audio_background, audio_eval):
+        embds_background = []
+        embds_eval = []
+        for sample in audio_background:
+            embd = self.fad.model.forward(sample.T.detach().numpy(), self.sr)
+            embds_background.append(embd.cpu().detach().numpy())
+        for sample in audio_eval:
+            embd = self.fad.model.forward(sample.T.detach().numpy(), self.sr)
+            embds_eval.append(embd.cpu().detach().numpy())
+        embds_background = np.concatenate(embds_background, axis=0)
+        embds_eval = np.concatenate(embds_eval, axis=0)
+        mu_background, sigma_background = self.fad.calculate_embd_statistics(
+            embds_background
+        )
+        mu_eval, sigma_eval = self.fad.calculate_embd_statistics(embds_eval)
+
+        fad_score = self.fad.calculate_frechet_distance(
+            mu_background, sigma_background, mu_eval, sigma_eval
+        )
+        return fad_score
+
+
 class RemFXModel(pl.LightningModule):
     def __init__(
         self,
@@ -35,7 +67,7 @@ class RemFXModel(pl.LightningModule):
             {
                 "SISDR": SISDRLoss(),
                 "STFT": STFTLoss(),
-                "L1": L1Loss(),
+                "FAD": FADLoss(sample_rate=sample_rate),
             }
         )
         # Log first batch metrics input vs output only once

shell_vars.sh

@@ -1,3 +1,3 @@
-export DATASET_ROOT="./data/GuitarSet"
+export DATASET_ROOT="./data/VocalSet"
 export WANDB_PROJECT="RemFX"
 export WANDB_ENTITY="mattricesound"