Allow for specifying max_effects as -1

README.md

@@ -28,8 +28,8 @@
 - `reverb`
 
 ## Train Main CLI Options
-- `max_kept_effects={n}` max number of <b> Kept </b> effects to apply to each file (default: 3)
-- `max_removed_effects={n}` max number of <b> Removed </b> effects to apply to each file (default: 4)
+- `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)
@@ -39,11 +39,18 @@
 - `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 max_removed_effects=4 shuffle_kept_effects=False shuffle_removed_effects=True accelerator='gpu' render_root='/home/username/datasets/vocalset'`
+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.
 
 ## Misc.
 By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
-
-

cfg/config.yaml

@@ -13,8 +13,8 @@ render_files: True
 render_root: "./data"
 accelerator: null
 
-max_kept_effects: 3
-max_removed_effects: 4
+max_kept_effects: -1
+max_removed_effects: -1
 shuffle_kept_effects: True
 shuffle_removed_effects: False
 effects_to_use:

remfx/datasets.py

@@ -10,6 +10,7 @@ from remfx import effects
 from tqdm import tqdm
 from remfx.utils import create_sequential_chunks
 import shutil
+from collections import OrderedSet
 
 
 # https://zenodo.org/record/1193957 -> VocalSet
@@ -22,12 +23,12 @@ class VocalSet(Dataset):
         self,
         root: str,
         sample_rate: int,
-        chunk_size: int = 3,
+        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 = 4,
-        max_removed_effects: int = 4,
+        max_kept_effects: int = -1,
+        max_removed_effects: int = 1,
         shuffle_kept_effects: bool = True,
         shuffle_removed_effects: bool = False,
         render_files: bool = True,
@@ -130,14 +131,25 @@ class VocalSet(Dataset):
                     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_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 effects: {kept_fx} (Up to {self.max_kept_effects}, chosen {kept_str}) -> Dry\n"
-            f"Apply effects: {rem_fx} (Up to {self.max_removed_effects}, chosen {rem_str}) -> Wet\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
 
@@ -151,7 +163,11 @@ class VocalSet(Dataset):
         else:
             effect_indices = torch.arange(len(self.effects_to_keep))
         # Up to max_kept_effects
-        effect_indices = effect_indices[: self.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]
@@ -168,6 +184,12 @@ class VocalSet(Dataset):
         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]