Clean project. Add 'all effects' to experiments

README.md

@@ -13,7 +13,7 @@
 4. Manually split singers into train, val, test directories
 
 ## Train model
-1. Change Wandb variables in `shell_vars.sh` and `source shell_vars.sh`
+1. Change Wandb and data root variables in `shell_vars.sh` and `source shell_vars.sh`
 2. `python scripts/train.py +exp=umx_distortion`
 or
 2. `python scripts/train.py +exp=demucs_distortion`
@@ -33,6 +33,12 @@ Ex. `python scripts/train.py +exp=umx_distortion trainer.accelerator='gpu' train
 - `compressor`
 - `distortion`
 - `reverb`
+- `all` (choose random effect to apply to each file)
 
 ## Misc.
-To skip rendering files, add `+datamodule.train_dataset.render_files=False +datamodule.val_dataset.render_files=False` to the command-line
+By default, files are rendered to `input_dir / processed / train/val/test`.
+To skip rendering files (use previously rendered), add `render_files=False` to the command-line
+
+Test
+Experiment dictates data, ckpt dictates model
+`python scripts/test.py +exp=umx_distortion.yaml +ckpt_path=test_ckpts/umx_dist.ckpt`

cfg/config.yaml

@@ -8,6 +8,7 @@ train: True
 sample_rate: 48000
 logs_dir: "./logs"
 log_every_n_steps: 1000
+render_files: True
 
 callbacks:
   model_checkpoint:
@@ -26,18 +27,27 @@ datamodule:
     _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
-    output_root: ${oc.env:OUTPUT_ROOT}/train
     chunk_size_in_sec: 6
     mode: "train"
     effect_types: ${effects.train_effects}
+    render_files: ${render_files}
   val_dataset:
     _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
-    output_root: ${oc.env:OUTPUT_ROOT}/val
     chunk_size_in_sec: 6
     mode: "val"
     effect_types: ${effects.val_effects}
+    render_files: ${render_files}
+  test_dataset:
+    _target_: remfx.datasets.VocalSet
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size_in_sec: 6
+    mode: "test"
+    effect_types: ${effects.val_effects}
+    render_files: ${render_files}
+
   batch_size: 16
   num_workers: 8
   pin_memory: True

cfg/config_guitarset.yaml

@@ -1,52 +0,0 @@
-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

cfg/config_guitfx.yaml

@@ -1,52 +0,0 @@
-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.GuitarFXDataset
-    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

cfg/effects/all.yaml

@@ -0,0 +1,70 @@
+# @package _global_
+effects:
+  train_effects:
+      Chorus:
+        _target_: remfx.effects.RandomPedalboardChorus
+        sample_rate: ${sample_rate}
+      Distortion:
+        _target_: remfx.effects.RandomPedalboardDistortion
+        sample_rate: ${sample_rate}
+        min_drive_db: -10
+        max_drive_db: 50
+      Compressor:
+        _target_: remfx.effects.RandomPedalboardCompressor
+        sample_rate: ${sample_rate}
+        min_threshold_db: -42.0
+        max_threshold_db: -20.0
+        min_ratio: 1.5
+        max_ratio: 6.0
+      Reverb:
+        _target_: remfx.effects.RandomPedalboardReverb
+        sample_rate: ${sample_rate}
+        min_room_size: 0.3
+        max_room_size: 1.0
+        min_damping: 0.2
+        max_damping: 1.0
+        min_wet_dry: 0.2
+        max_wet_dry: 0.8
+        min_width: 0.2
+        max_width: 1.0
+  val_effects:
+      Chorus:
+        _target_: remfx.effects.RandomPedalboardChorus
+        sample_rate: ${sample_rate}
+        min_rate_hz: 1.0
+        max_rate_hz: 1.0
+        min_depth: 0.3
+        max_depth: 0.3
+        min_centre_delay_ms: 7.5
+        max_centre_delay_ms: 7.5
+        min_feedback: 0.4
+        max_feedback: 0.4
+        min_mix: 0.4
+        max_mix: 0.4
+      Distortion:
+        _target_: remfx.effects.RandomPedalboardDistortion
+        sample_rate: ${sample_rate}
+        min_drive_db: 30
+        max_drive_db: 30
+      Compressor:
+        _target_: remfx.effects.RandomPedalboardCompressor
+        sample_rate: ${sample_rate}
+        min_threshold_db: -32
+        max_threshold_db: -32
+        min_ratio: 3.0
+        max_ratio: 3.0
+        min_attack_ms: 10.0
+        max_attack_ms: 10.0
+        min_release_ms: 40.0
+        max_release_ms: 40.0
+      Reverb:
+        _target_: remfx.effects.RandomPedalboardReverb
+        sample_rate: ${sample_rate}
+        min_room_size: 0.5
+        max_room_size: 0.5
+        min_damping: 0.5
+        max_damping: 0.5
+        min_wet_dry: 0.4
+        max_wet_dry: 0.4
+        min_width: 0.5
+        max_width: 0.5

cfg/exp/demucs_all.yaml

@@ -0,0 +1,4 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all

cfg/exp/umx_all.yaml

@@ -0,0 +1,4 @@
+# @package _global_
+defaults:
+  - override /model: umx
+  - override /effects: all

remfx/datasets.py

@@ -1,179 +1,16 @@
 import torch
-from torch.utils.data import Dataset, DataLoader, random_split
+from torch.utils.data import Dataset, DataLoader
 import torchaudio
-import torchaudio.transforms as T
 import torch.nn.functional as F
 from pathlib import Path
 import pytorch_lightning as pl
-from typing import Any, List, Tuple
+from typing import Any, List
 from remfx import effects
-from pedalboard import (
-    Pedalboard,
-    Chorus,
-    Reverb,
-    Compressor,
-    Phaser,
-    Delay,
-    Distortion,
-    Limiter,
-)
 from tqdm import tqdm
+from remfx.utils import create_sequential_chunks
 
-# https://zenodo.org/record/7044411/ -> GuitarFX
-# https://zenodo.org/record/3371780  -> GuitarSet
 # https://zenodo.org/record/1193957 -> VocalSet
 
-deterministic_effects = {
-    "Distortion": Pedalboard([Distortion()]),
-    "Compressor": Pedalboard([Compressor()]),
-    "Chorus": Pedalboard([Chorus()]),
-    "Phaser": Pedalboard([Phaser()]),
-    "Delay": Pedalboard([Delay()]),
-    "Reverb": Pedalboard([Reverb()]),
-    "Limiter": Pedalboard([Limiter()]),
-}
-
-
-class GuitarFXDataset(Dataset):
-    def __init__(
-        self,
-        root: str,
-        sample_rate: int,
-        chunk_size_in_sec: int = 3,
-        effect_types: List[str] = None,
-    ):
-        super().__init__()
-        self.wet_files = []
-        self.dry_files = []
-        self.chunks = []
-        self.labels = []
-        self.song_idx = []
-        self.root = Path(root)
-        self.chunk_size_in_sec = chunk_size_in_sec
-        self.sample_rate = sample_rate
-
-        if effect_types is None:
-            effect_types = [
-                d.name for d in self.root.iterdir() if d.is_dir() and d != "Clean"
-            ]
-        current_file = 0
-        for i, effect in enumerate(effect_types):
-            for pickup in Path(self.root / effect).iterdir():
-                wet_files = sorted(list(pickup.glob("*.wav")))
-                dry_files = sorted(
-                    list(self.root.glob(f"Clean/{pickup.name}/**/*.wav"))
-                )
-                self.wet_files += wet_files
-                self.dry_files += dry_files
-                self.labels += [i] * len(wet_files)
-                for audio_file in wet_files:
-                    chunk_starts, orig_sr = create_sequential_chunks(
-                        audio_file, self.chunk_size_in_sec
-                    )
-                    self.chunks += chunk_starts
-                    self.song_idx += [current_file] * len(chunk_starts)
-                    current_file += 1
-        print(
-            f"Found {len(self.wet_files)} wet files and {len(self.dry_files)} dry files.\n"
-            f"Total chunks: {len(self.chunks)}"
-        )
-        self.resampler = T.Resample(orig_sr, sample_rate)
-
-    def __len__(self):
-        return len(self.chunks)
-
-    def __getitem__(self, idx):
-        # Load effected and "clean" audio
-        song_idx = self.song_idx[idx]
-        x, sr = torchaudio.load(self.wet_files[song_idx])
-        y, sr = torchaudio.load(self.dry_files[song_idx])
-        effect_label = self.labels[song_idx]  # Effect label
-
-        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]
-        y = y[:, chunk_start : chunk_start + chunk_size_in_samples]
-
-        resampled_x = self.resampler(x)
-        resampled_y = self.resampler(y)
-        # 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])
-            )
-        if resampled_y.shape[-1] < chunk_size_in_samples:
-            resampled_y = F.pad(
-                resampled_y, (0, chunk_size_in_samples - resampled_y.shape[1])
-            )
-        return (resampled_x, resampled_y, effect_label)
-
-
-class GuitarSet(Dataset):
-    def __init__(
-        self,
-        root: str,
-        sample_rate: int,
-        chunk_size_in_sec: int = 3,
-        effect_types: List[torch.nn.Module] = None,
-    ):
-        super().__init__()
-        self.chunks = []
-        self.song_idx = []
-        self.root = Path(root)
-        self.chunk_size_in_sec = chunk_size_in_sec
-        self.files = sorted(list(self.root.glob("./**/*.wav")))
-        self.sample_rate = sample_rate
-        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)
-        self.mode = "train"
-
-    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)
-
 
 class VocalSet(Dataset):
     def __init__(
@@ -183,7 +20,6 @@ class VocalSet(Dataset):
         chunk_size_in_sec: int = 3,
         effect_types: List[torch.nn.Module] = None,
         render_files: bool = True,
-        output_root: str = "processed",
         mode: str = "train",
     ):
         super().__init__()
@@ -199,14 +35,15 @@ class VocalSet(Dataset):
         self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
         self.effect_types = effect_types
 
-        self.output_root = Path(output_root)
+        self.processed_root = self.root / "processed" / self.mode
 
         self.num_chunks = 0
         print("Total files:", len(self.files))
         print("Processing files...")
         if render_files:
-            self.output_root.mkdir(parents=True, exist_ok=True)
-            for i, audio_file in tqdm(enumerate(self.files)):
+            # Split audio file into chunks, resample, then apply random effects
+            self.processed_root.mkdir(parents=True, exist_ok=True)
+            for audio_file in tqdm(self.files, total=len(self.files)):
                 chunks, orig_sr = create_sequential_chunks(
                     audio_file, self.chunk_size_in_sec
                 )
@@ -220,14 +57,16 @@ class VocalSet(Dataset):
                             resampled_chunk,
                             (0, chunk_size_in_samples - resampled_chunk.shape[1]),
                         )
+                    # Apply effect
                     effect_idx = torch.rand(1).item() * len(self.effect_types.keys())
                     effect_name = list(self.effect_types.keys())[int(effect_idx)]
                     effect = self.effect_types[effect_name]
                     effected_input = effect(resampled_chunk)
+                    # Normalize
                     normalized_input = self.normalize(effected_input)
                     normalized_target = self.normalize(resampled_chunk)
 
-                    output_dir = self.output_root / str(self.num_chunks)
+                    output_dir = self.processed_root / str(self.num_chunks)
                     output_dir.mkdir(exist_ok=True)
                     torchaudio.save(
                         output_dir / "input.wav", normalized_input, self.sample_rate
@@ -235,9 +74,10 @@ class VocalSet(Dataset):
                     torchaudio.save(
                         output_dir / "target.wav", normalized_target, self.sample_rate
                     )
+                    torch.save(effect_name, output_dir / "effect_name.pt")
                     self.num_chunks += 1
         else:
-            self.num_chunks = len(list(self.output_root.glob("./**/*.wav")))
+            self.num_chunks = len(list(self.processed_root.iterdir()))
 
         print(
             f"Found {len(self.files)} {self.mode} files .\n"
@@ -248,95 +88,12 @@ class VocalSet(Dataset):
         return self.num_chunks
 
     def __getitem__(self, idx):
-        # Load audio
-        input_file = self.output_root / str(idx) / "input.wav"
-        target_file = self.output_root / str(idx) / "target.wav"
+        input_file = self.processed_root / str(idx) / "input.wav"
+        target_file = self.processed_root / str(idx) / "target.wav"
+        effect_name = torch.load(self.processed_root / str(idx) / "effect_name.pt")
         input, sr = torchaudio.load(input_file)
         target, sr = torchaudio.load(target_file)
-        return (input, target, "")
-
-
-def create_random_chunks(
-    audio_file: str, chunk_size: int, num_chunks: int
-) -> Tuple[List[Tuple[int, int]], int]:
-    """Create num_chunks random chunks of size chunk_size (seconds)
-    from an audio file.
-    Return sample_index of start of each chunk and original sr
-    """
-    audio, sr = torchaudio.load(audio_file)
-    chunk_size_in_samples = chunk_size * sr
-    if chunk_size_in_samples >= audio.shape[-1]:
-        chunk_size_in_samples = audio.shape[-1] - 1
-    chunks = []
-    for i in range(num_chunks):
-        start = torch.randint(0, audio.shape[-1] - chunk_size_in_samples, (1,)).item()
-        chunks.append(start)
-    return chunks, sr
-
-
-def create_sequential_chunks(
-    audio_file: str, chunk_size: int
-) -> Tuple[List[Tuple[int, int]], int]:
-    """Create sequential chunks of size chunk_size (seconds) from an audio file.
-    Return sample_index of start of each chunk and original sr
-    """
-    chunks = []
-    audio, sr = torchaudio.load(audio_file)
-    chunk_size_in_samples = chunk_size * sr
-    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size_in_samples)
-    for start in chunk_starts:
-        if start + chunk_size_in_samples > audio.shape[-1]:
-            break
-        chunks.append(audio[:, start : start + chunk_size_in_samples])
-    return chunks, sr
-
-
-class Datamodule(pl.LightningDataModule):
-    def __init__(
-        self,
-        dataset,
-        *,
-        val_split: float,
-        batch_size: int,
-        num_workers: int,
-        pin_memory: bool = False,
-        **kwargs: int,
-    ) -> None:
-        super().__init__()
-        self.dataset = dataset
-        self.val_split = val_split
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-        self.data_train: Any = None
-        self.data_val: Any = None
-
-    def setup(self, stage: Any = None) -> None:
-        split = [1.0 - self.val_split, self.val_split]
-        train_size = round(split[0] * len(self.dataset))
-        val_size = round(split[1] * len(self.dataset))
-        self.data_train, self.data_val = random_split(
-            self.dataset, [train_size, val_size]
-        )
-        self.data_val.dataset.mode = "val"
-
-    def train_dataloader(self) -> DataLoader:
-        return DataLoader(
-            dataset=self.data_train,
-            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.data_val,
-            batch_size=self.batch_size,
-            num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
-            shuffle=False,
-        )
+        return (input, target, effect_name)
 
 
 class VocalSetDatamodule(pl.LightningDataModule):
@@ -344,6 +101,7 @@ class VocalSetDatamodule(pl.LightningDataModule):
         self,
         train_dataset,
         val_dataset,
+        test_dataset,
         *,
         batch_size: int,
         num_workers: int,
@@ -353,6 +111,7 @@ class VocalSetDatamodule(pl.LightningDataModule):
         super().__init__()
         self.train_dataset = train_dataset
         self.val_dataset = val_dataset
+        self.test_dataset = test_dataset
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.pin_memory = pin_memory
@@ -377,3 +136,12 @@ class VocalSetDatamodule(pl.LightningDataModule):
             pin_memory=self.pin_memory,
             shuffle=False,
         )
+
+    def test_dataloader(self) -> DataLoader:
+        return DataLoader(
+            dataset=self.test_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+        )

remfx/models.py

@@ -7,44 +7,12 @@ 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 remfx.utils import FADLoss
 
 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.fad.model = self.fad.model.to("cpu")
-        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.cpu().detach().numpy(), self.sr)
-            embds_background.append(embd.cpu().detach().numpy())
-        for sample in audio_eval:
-            embd = self.fad.model.forward(sample.T.cpu().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,
@@ -97,6 +65,10 @@ class RemFXModel(pl.LightningModule):
         loss = self.common_step(batch, batch_idx, mode="valid")
         return loss
 
+    def test_step(self, batch, batch_idx):
+        loss = self.common_step(batch, batch_idx, mode="test")
+        return loss
+
     def common_step(self, batch, batch_idx, mode: str = "train"):
         loss, output = self.model(batch)
         self.log(f"{mode}_loss", loss)
@@ -121,6 +93,7 @@ class RemFXModel(pl.LightningModule):
         return loss
 
     def on_train_batch_start(self, batch, batch_idx):
+        # Log initial audio
         if self.log_train_audio:
             x, y, label = batch
             # Concat samples together for easier viewing in dashboard
@@ -143,48 +116,47 @@ class RemFXModel(pl.LightningModule):
             )
             self.log_train_audio = False
 
-    def on_validation_epoch_start(self):
-        self.log_next = True
-
     def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
-        if self.log_next:
-            x, target, label = batch
-            # Log Input Metrics
-            for metric in self.metrics:
-                # SISDR returns negative values, so negate them
-                if metric == "SISDR":
-                    negate = -1
-                else:
-                    negate = 1
-                self.log(
-                    f"Input_{metric}",
-                    negate * self.metrics[metric](x, target),
-                    on_step=False,
-                    on_epoch=True,
-                    logger=True,
-                    prog_bar=True,
-                    sync_dist=True,
-                )
-
-            self.model.eval()
-            with torch.no_grad():
-                y = self.model.sample(x)
+        x, target, label = batch
+        # Log Input Metrics
+        for metric in self.metrics:
+            # SISDR returns negative values, so negate them
+            if metric == "SISDR":
+                negate = -1
+            else:
+                negate = 1
+            self.log(
+                f"Input_{metric}",
+                negate * self.metrics[metric](x, target),
+                on_step=False,
+                on_epoch=True,
+                logger=True,
+                prog_bar=True,
+                sync_dist=True,
+            )
 
-            # Concat samples together for easier viewing in dashboard
-            # 2 seconds of silence between each sample
-            silence = torch.zeros_like(x)
-            silence = silence[:, : self.sample_rate * 2]
+        self.model.eval()
+        with torch.no_grad():
+            y = self.model.sample(x)
+
+        # Concat samples together for easier viewing in dashboard
+        # 2 seconds of silence between each sample
+        silence = torch.zeros_like(x)
+        silence = silence[:, : self.sample_rate * 2]
+
+        concat_samples = torch.cat([y, silence, x, silence, target], dim=-1)
+        log_wandb_audio_batch(
+            logger=self.logger,
+            id="prediction_input_target",
+            samples=concat_samples.cpu(),
+            sampling_rate=self.sample_rate,
+            caption=f"Epoch {self.current_epoch}",
+        )
+        self.log_next = False
+        self.model.train()
 
-            concat_samples = torch.cat([y, silence, x, silence, target], dim=-1)
-            log_wandb_audio_batch(
-                logger=self.logger,
-                id="prediction_input_target",
-                samples=concat_samples.cpu(),
-                sampling_rate=self.sample_rate,
-                caption=f"Epoch {self.current_epoch}",
-            )
-            self.log_next = False
-            self.model.train()
+    def on_test_batch_start(self, batch, batch_idx, dataloader_idx):
+        return self.on_validation_batch_start(batch, batch_idx, dataloader_idx)
 
 
 class OpenUnmixModel(torch.nn.Module):

remfx/utils.py

@@ -1,8 +1,12 @@
 import logging
-from typing import List
+from typing import List, Tuple
 import pytorch_lightning as pl
 from omegaconf import DictConfig
 from pytorch_lightning.utilities import rank_zero_only
+from frechet_audio_distance import FrechetAudioDistance
+import numpy as np
+import torch
+import torchaudio
 
 
 def get_logger(name=__name__) -> logging.Logger:
@@ -69,3 +73,69 @@ def log_hyperparameters(
         hparams["callbacks"] = config["callbacks"]
 
     logger.experiment.config.update(hparams)
+
+
+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.fad.model = self.fad.model.to("cpu")
+        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.cpu().detach().numpy(), self.sr)
+            embds_background.append(embd.cpu().detach().numpy())
+        for sample in audio_eval:
+            embd = self.fad.model.forward(sample.T.cpu().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
+
+
+def create_random_chunks(
+    audio_file: str, chunk_size: int, num_chunks: int
+) -> Tuple[List[Tuple[int, int]], int]:
+    """Create num_chunks random chunks of size chunk_size (seconds)
+    from an audio file.
+    Return sample_index of start of each chunk and original sr
+    """
+    audio, sr = torchaudio.load(audio_file)
+    chunk_size_in_samples = chunk_size * sr
+    if chunk_size_in_samples >= audio.shape[-1]:
+        chunk_size_in_samples = audio.shape[-1] - 1
+    chunks = []
+    for i in range(num_chunks):
+        start = torch.randint(0, audio.shape[-1] - chunk_size_in_samples, (1,)).item()
+        chunks.append(start)
+    return chunks, sr
+
+
+def create_sequential_chunks(
+    audio_file: str, chunk_size: int
+) -> Tuple[List[Tuple[int, int]], int]:
+    """Create sequential chunks of size chunk_size (seconds) from an audio file.
+    Return sample_index of start of each chunk and original sr
+    """
+    chunks = []
+    audio, sr = torchaudio.load(audio_file)
+    chunk_size_in_samples = chunk_size * sr
+    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size_in_samples)
+    for start in chunk_starts:
+        if start + chunk_size_in_samples > audio.shape[-1]:
+            break
+        chunks.append(audio[:, start : start + chunk_size_in_samples])
+    return chunks, sr

shell_vars.sh

@@ -1,4 +1,3 @@
 export DATASET_ROOT="./data/VocalSet"
-export OUTPUT_ROOT="/scratch/VocalSet/processed"
 export WANDB_PROJECT="RemFX"
 export WANDB_ENTITY="mattricesound"