Merge pull request #28 from mhrice/refactor-config

.gitignore

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

README.md

@@ -1,4 +1,6 @@
 
+# Setup
+
 ## Install Packages
 1. `python3 -m venv env`
 2. `source env/bin/activate`
@@ -12,34 +14,45 @@
 3. `unzip VocalSet.zip`
 4. Manually split singers into train, val, test directories
 
-## Train model
+# Training
+## Steps
 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.
-
-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
+2. `python scripts/train.py model=demucs "effects_to_remove=[distortion]"`
+
+## Main CLI Options
+- `max_kept_effects={n}` max number of <b> Kept </b> effects to apply to each file. Set to -1 to always use all effects (default: -1)
+- `max_removed_effects={n}` max number of <b> Removed </b> effects to apply to each file. Set to -1 to always use all effects (default: -1)
+- `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)
+- `accelerator=null/gpu` Use GPU (1 device) (default: False)
+- `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)
+
+Note that "kept effects" are calculated from the difference between `effects_to_use` and `effects_to_remove`.
+
+Example: `python scripts/train.py model=demucs "effects_to_use=[distortion, reverb, chorus]" "effects_to_remove=[distortion]" max_kept_effects=2 max_removed_effects=4 shuffle_kept_effects=False shuffle_removed_effects=True accelerator='gpu' render_root=/scratch/VocalSet'`
+
+Printout:
+```
+Effect Summary:
+Apply kept effects: ['chorus', 'reverb'] (Up to 2, chosen in order) -> Dry
+Apply remove effects: ['distortion'] (Up to 4, chosen randomly) -> Wet
+```
+
+See `cfg/config.yaml` for more options that can be specified on the command line.
+
+## 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`
 
 ## 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
@@ -12,6 +11,22 @@ chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
 render_root: "./data"
+accelerator: null
+
+max_kept_effects: -1
+max_removed_effects: -1
+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:
@@ -35,9 +50,13 @@ 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}
+    max_removed_effects: ${max_removed_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 +65,13 @@ 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}
+    max_removed_effects: ${max_removed_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 +80,13 @@ 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}
+    max_removed_effects: ${max_removed_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
     render_files: ${render_files}
     render_root: ${render_root}
 
@@ -85,7 +112,8 @@ trainer:
   enable_model_summary: False
   log_every_n_steps: 1 # Logs metrics every N batches
   accumulate_grad_batches: 1
-  accelerator: null
+  accelerator: ${accelerator}
   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,11 +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
+from ordered_set import OrderedSet
+
 
 # https://zenodo.org/record/1193957 -> VocalSet
 
@@ -21,10 +23,14 @@ class VocalSet(Dataset):
         self,
         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,
+        chunk_size: int = 262144,
+        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,
+        max_removed_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 +43,20 @@ 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.max_removed_effects = max_removed_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 +112,96 @@ 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(
+            OrderedSet(self.effects_to_use) - OrderedSet(self.effects_to_remove)
+        )
+        kept_str = "randomly" if self.shuffle_kept_effects else "in order"
+        rem_fx = self.effects_to_remove
+        rem_str = "randomly" if self.shuffle_removed_effects else "in order"
+        if self.max_kept_effects == -1:
+            num_kept_str = len(kept_fx)
+        else:
+            num_kept_str = f"Up to {self.max_kept_effects}"
+        if self.max_removed_effects == -1:
+            num_rem_str = len(rem_fx)
+        else:
+            num_rem_str = f"Up to {self.max_removed_effects}"
+
+        print(
+            f"Effect Summary: \n"
+            f"Apply kept effects: {kept_fx} ({num_kept_str}, chosen {kept_str}) -> Dry\n"
+            f"Apply remove effects: {rem_fx} ({num_rem_str}, chosen {rem_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
+        # Shuffle effects if specified
+        if self.shuffle_kept_effects:
+            effect_indices = torch.randperm(len(self.effects_to_keep))
+        else:
+            effect_indices = torch.arange(len(self.effects_to_keep))
+        # Up to max_kept_effects
+        if self.max_kept_effects != -1:
+            num_kept_effects = int(torch.rand(1).item() * (self.max_kept_effects)) + 1
+        else:
+            num_kept_effects = len(self.effects_to_keep)
+        effect_indices = effect_indices[:num_kept_effects]
+        # Index in effect settings
+        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]
+        # Apply
+        for effect in effects_to_apply:
+            dry = effect(dry)
+            labels.append(ALL_EFFECTS.index(type(effect)))
+
+        # Apply effects_to_remove
+        # Shuffle effects if specified
         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))
+        # Up to max_removed_effects
+        if self.max_removed_effects != -1:
+            num_kept_effects = (
+                int(torch.rand(1).item() * (self.max_removed_effects)) + 1
+            )
+        else:
+            num_kept_effects = len(self.effects_to_remove)
+        effect_indices = effect_indices[: self.max_removed_effects]
+        # Index in effect settings
+        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]
+        # 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)

setup.py

@@ -47,6 +47,7 @@ setup(
         "pyloudnorm",
         "pedalboard",
         "frechet_audio_distance",
+        "ordered-set",
     ],
     include_package_data=True,
     license="Apache License 2.0",