Refactor commit, see README for details

.gitignore

@@ -8,4 +8,5 @@ __pycache__/
 lightning_logs/
 outputs/
 logs/
-.vscode/
+.vscode/
+ckpts/

README.md

@@ -14,32 +14,35 @@
 
 ## Train model
 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`
-See cfg for more options. Generally they are `+exp={model}_{effect}`
-Models and effects detailed below.
+2. `python scripts/train.py model=demucs "effects_to_remove=[distortion]"`
 
-To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-line
 
-Ex. `python scripts/train.py +exp=umx_distortion trainer.accelerator='gpu' trainer.devices=1`
-
-### Current Models
+## Models
 - `umx`
 - `demucs`
 
-### Current Effects
+## Effects
 - `chorus`
 - `compressor`
 - `distortion`
 - `reverb`
-- `all` (choose random effect to apply to each file)
 
-### Testing
-Experiment dictates data, ckpt dictates model
-`python scripts/test.py +exp=umx_distortion.yaml +ckpt_path=test_ckpts/umx_dist.ckpt`
+## Train CLI Options
+- `max_kept_effects={n}` max number of <b> Kept </b> effects to apply to each file (default: 3)
+- `model={model}` architecture to use (see 'Models')
+- `shuffle_kept_effects=True/False` Shuffle kept effects (default: True)
+- `shuffle_removed_effects=True/False` Shuffle Removed effects (default: False)
+- `effects_to_use={effect}` Effects to use (see 'Effects') (default: all in the list)
+- `effects_to_remove={effect}` Effects to remove (see 'Effects') (default: all in the list)
+- `trainer.accelerator='gpu'` : Use GPU (default: None)
+- `trainer.devices={n}` Number of GPUs to use (default: 1)
+- `render_files=True/False` Render files. Disable to skip rendering stage (default: True)
+- `render_root={path/to/dir}`. Root directory to render files to (default: DATASET_ROOT)
+
+Example: `python scripts/train.py model=demucs "effects_to_use=[distortion, reverb]" "effects_to_remove=[distortion]" "max_kept_effects=2" "shuffle_kept_effects=False" "shuffle_removed_effects=True" trainer.accelerator='gpu' trainer.devices=2`
+
 
 ## Misc.
 By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
-To skip rendering files (use previously rendered), add `render_files=False` to the command-line (added to test by default).
-To change the rendered location, add `render_root={path/to/dir}` to the command-line (use this for train and test)
+
+

cfg/applied_effects/all.yaml

@@ -1,31 +0,0 @@
-# @package _global_
-applied_effects:
-  Chorus:
-    _target_: remfx.effects.RandomPedalboardChorus
-    sample_rate: ${sample_rate}
-    min_depth: 0.2
-    min_mix: 0.3
-  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
-

cfg/config.yaml

@@ -1,10 +1,9 @@
 defaults:
   - _self_
   - model: null
-  - applied_effects: null
-  - effect_to_remove: null
+  - effects: all
+
 
-max_effects_per_file: 3
 seed: 12345
 train: True
 sample_rate: 48000
@@ -13,6 +12,20 @@ logs_dir: "./logs"
 render_files: True
 render_root: "./data"
 
+max_kept_effects: 3
+shuffle_kept_effects: True
+shuffle_removed_effects: False
+effects_to_use:
+  - compressor
+  - distortion
+  - reverb
+  - chorus
+effects_to_remove:
+  - compressor
+  - distortion
+  - reverb
+  - chorus
+
 callbacks:
   model_checkpoint:
     _target_: pytorch_lightning.callbacks.ModelCheckpoint
@@ -35,9 +48,12 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "train"
-    applied_effects: ${applied_effects}
-    effect_to_remove: ${effect_to_remove}
-    max_effects_per_file: ${max_effects_per_file}
+    effect_modules: ${effects}
+    effects_to_use: ${effects_to_use}
+    effects_to_remove: ${effects_to_remove}
+    max_kept_effects: ${max_kept_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
     render_files: ${render_files}
     render_root: ${render_root}
   val_dataset:
@@ -46,9 +62,12 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "val"
-    applied_effects: ${applied_effects}
-    effect_to_remove: ${effect_to_remove}
-    max_effects_per_file: ${max_effects_per_file}
+    effect_modules: ${effects}
+    effects_to_use: ${effects_to_use}
+    effects_to_remove: ${effects_to_remove}
+    max_kept_effects: ${max_kept_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
     render_files: ${render_files}
     render_root: ${render_root}
   test_dataset:
@@ -57,9 +76,12 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "test"
-    applied_effects: ${applied_effects}
-    effect_to_remove: ${effect_to_remove}
-    max_effects_per_file: ${max_effects_per_file}
+    effect_modules: ${effects}
+    effects_to_use: ${effects_to_use}
+    effects_to_remove: ${effects_to_remove}
+    max_kept_effects: ${max_kept_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
     render_files: ${render_files}
     render_root: ${render_root}
 
@@ -89,3 +111,4 @@ trainer:
   devices: 1
   gradient_clip_val: 10.0
   max_steps: 50000
+

cfg/effect_to_remove/all.yaml

@@ -1,31 +0,0 @@
-# @package _global_
-effect_to_remove:
-  Chorus:
-    _target_: remfx.effects.RandomPedalboardChorus
-    sample_rate: ${sample_rate}
-    min_depth: 0.2
-    min_mix: 0.3
-  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
-

cfg/effect_to_remove/chorus.yaml

@@ -1,7 +0,0 @@
-# @package _global_
-effect_to_remove:
-  Chorus:
-    _target_: remfx.effects.RandomPedalboardChorus
-    sample_rate: ${sample_rate}
-    min_depth: 0.2
-    min_mix: 0.3

cfg/effect_to_remove/compressor.yaml

@@ -1,9 +0,0 @@
-# @package _global_
-effect_to_remove:
-  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

cfg/effect_to_remove/distortion.yaml

@@ -1,7 +0,0 @@
-# @package _global_
-effect_to_remove:
-  Distortion:
-    _target_: remfx.effects.RandomPedalboardDistortion
-    sample_rate: ${sample_rate}
-    min_drive_db: 10
-    max_drive_db: 50

cfg/effect_to_remove/reverb.yaml

@@ -1,13 +0,0 @@
-# @package _global_
-effect_to_remove:
-  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

cfg/effects/all.yaml

@@ -0,0 +1,31 @@
+# @package _global_
+
+effects:
+  chorus:
+      _target_: remfx.effects.RandomPedalboardChorus
+      sample_rate: ${sample_rate}
+      min_depth: 0.2
+      min_mix: 0.3
+  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

cfg/exp/demucs_all.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: demucs
-  - override /applied_effects: all
-  - override /effect_to_remove: all

cfg/exp/demucs_chorus.yaml

@@ -1,6 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: demucs
-  - override /applied_effects: all
-  - override /effect_to_remove: chorus
-

cfg/exp/demucs_compressor.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: demucs
-  - override /applied_effects: all
-  - override /effect_to_remove: distortion

cfg/exp/demucs_distortion.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: demucs
-  - override /applied_effects: all
-  - override /effect_to_remove: distortion

cfg/exp/demucs_reverb.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: demucs
-  - override /applied_effects: all
-  - override /effect_to_remove: reverb

cfg/exp/umx_all.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: umx
-  - override /applied_effects: all
-  - override /effect_to_remove: all

cfg/exp/umx_chorus.yaml

@@ -1,6 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: umx
-  - override /applied_effects: all
-  - override /effect_to_remove: chorus
-

cfg/exp/umx_compressor.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: umx
-  - override /applied_effects: all
-  - override /effect_to_remove: distortion

cfg/exp/umx_distortion.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: umx
-  - override /applied_effects: all
-  - override /effect_to_remove: distortion

cfg/exp/umx_reverb.yaml

@@ -1,5 +0,0 @@
-# @package _global_
-defaults:
-  - override /model: umx
-  - override /applied_effects: all
-  - override /effect_to_remove: reverb

remfx/datasets.py

@@ -5,12 +5,13 @@ import torchaudio
 from pathlib import Path
 import pytorch_lightning as pl
 import sys
-from typing import Any, Dict
+from typing import Any, List, Dict
 from remfx import effects
 from tqdm import tqdm
 from remfx.utils import create_sequential_chunks
 import shutil
 
+
 # https://zenodo.org/record/1193957 -> VocalSet
 
 ALL_EFFECTS = effects.Pedalboard_Effects
@@ -22,9 +23,12 @@ class VocalSet(Dataset):
         root: str,
         sample_rate: int,
         chunk_size: int = 3,
-        applied_effects: Dict[str, torch.nn.Module] = None,
-        effect_to_remove: Dict[str, torch.nn.Module] = None,
-        max_effects_per_file: int = 1,
+        effect_modules: List[Dict[str, torch.nn.Module]] = None,
+        effects_to_use: List[str] = None,
+        effects_to_remove: List[str] = None,
+        max_kept_effects: int = 1,
+        shuffle_kept_effects: bool = True,
+        shuffle_removed_effects: bool = False,
         render_files: bool = True,
         render_root: str = None,
         mode: str = "train",
@@ -37,17 +41,19 @@ class VocalSet(Dataset):
         self.chunk_size = chunk_size
         self.sample_rate = sample_rate
         self.mode = mode
-        self.max_effects_per_file = max_effects_per_file
-        self.effect_to_remove = effect_to_remove
         mode_path = self.root / self.mode
         self.files = sorted(list(mode_path.glob("./**/*.wav")))
+        self.max_kept_effects = max_kept_effects
+        self.effects_to_use = effects_to_use
+        self.effects_to_remove = effects_to_remove
         self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
-        self.applied_effects = applied_effects
-        self.effect_to_remove_name = "_".join([e for e in self.effect_to_remove])
+        self.effects = effect_modules
+        self.shuffle_kept_effects = shuffle_kept_effects
+        self.shuffle_removed_effects = shuffle_removed_effects
 
-        effect_str = "__".join([e for e in self.applied_effects])
-        effect_str += f"_{self.effect_to_remove_name}"
-        self.proc_root = self.render_root / "processed" / effect_str / self.mode
+        effects_string = "_".join(self.effects_to_use + ["_"] + self.effects_to_remove)
+        self.effects_to_keep = self.validate_effect_input()
+        self.proc_root = self.render_root / "processed" / effects_string / self.mode
 
         if self.proc_root.exists() and len(list(self.proc_root.iterdir())) > 0:
             print("Found processed files.")
@@ -103,38 +109,66 @@ class VocalSet(Dataset):
         target, sr = torchaudio.load(target_file)
         return (input, target, effect_name)
 
+    def validate_effect_input(self):
+        for effect in self.effects.values():
+            if type(effect) not in ALL_EFFECTS:
+                raise ValueError(
+                    f"Effect {effect} not found in ALL_EFFECTS. "
+                    f"Please choose from {ALL_EFFECTS}"
+                )
+        for effect in self.effects_to_use:
+            if effect not in self.effects.keys():
+                raise ValueError(
+                    f"Effect {effect} not found in self.effects. "
+                    f"Please choose from {self.effects.keys()}"
+                )
+        for effect in self.effects_to_remove:
+            if effect not in self.effects.keys():
+                raise ValueError(
+                    f"Effect {effect} not found in self.effects. "
+                    f"Please choose from {self.effects.keys()}"
+                )
+        kept_fx = list(set(self.effects_to_use) - set(self.effects_to_remove))
+        kept_str = "randomly" if self.shuffle_kept_effects else "in order"
+        removed_str = "randomly" if self.shuffle_removed_effects else "in order"
+        rem_fx = self.effects_to_remove
+        print(
+            f"Effect Summary: \n"
+            f"Apply effects: {kept_fx} (Up to {self.max_kept_effects}, chosen {kept_str}) -> Dry\n"
+            f"Apply effects: {rem_fx} (All {len(rem_fx)}, chosen {removed_str}) -> Wet\n"
+        )
+        return kept_fx
+
     def process_effects(self, dry: torch.Tensor):
-        # Apply random number of effects up to num_effects - 1 (excluding effect_to_remove)
-        if self.max_effects_per_file > 1:
-            num_effects = torch.randint(self.max_effects_per_file - 1, (1,)).item()
-            # Remove effect to remove from applied effects if present
-            for effect in self.effect_to_remove:
-                self.applied_effects.pop(effect, None)
-
-            # Choose random effects to apply
-            effect_indices = torch.randperm(len(self.applied_effects.keys()))[
-                :num_effects
-            ]
-            effects_to_apply = [
-                list(self.applied_effects.keys())[i] for i in effect_indices
-            ]
-            labels = []
-            for effect_name in effects_to_apply:
-                effect = self.applied_effects[effect_name]
-                dry = effect(dry)
-                labels.append(ALL_EFFECTS.index(type(effect)))
-
-        # Apply effect_to_remove
+        labels = []
+
+        # Apply Kept Effects
+        if self.shuffle_kept_effects:
+            effect_indices = torch.randperm(len(self.effects_to_keep))
+        else:
+            effect_indices = torch.arange(len(self.effects_to_keep))
+        effect_names_to_apply = [self.effects_to_keep[i] for i in effect_indices]
+        effects_to_apply = [self.effects[i] for i in effect_names_to_apply]
+        for effect in effects_to_apply:
+            dry = effect(dry)
+            labels.append(ALL_EFFECTS.index(type(effect)))
+        print(labels)
+
+        # Apply effects_to_remove
         wet = torch.clone(dry)
-        for effect_name in self.effect_to_remove:
-            effect = self.effect_to_remove[effect_name]
-            wet = effect(dry)
+        if self.shuffle_removed_effects:
+            effect_indices = torch.randperm(len(self.effects_to_remove))
+        else:
+            effect_indices = torch.arange(len(self.effects_to_remove))
+        effect_names_to_apply = [self.effects_to_remove[i] for i in effect_indices]
+        effects_to_apply = [self.effects[i] for i in effect_names_to_apply]
+        for effect in effects_to_apply:
+            wet = effect(wet)
             labels.append(ALL_EFFECTS.index(type(effect)))
 
         # Convert labels to one-hot
         one_hot = F.one_hot(torch.tensor(labels), num_classes=len(ALL_EFFECTS))
         effects_present = torch.sum(one_hot, dim=0).float()
-
         # Normalize
         normalized_dry = self.normalize(dry)
         normalized_wet = self.normalize(wet)