Merge pull request #24 from mhrice/new-metrics

README.md

@@ -22,7 +22,7 @@ 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`
+Ex. `python scripts/train.py +exp=umx_distortion trainer.accelerator='gpu' trainer.devices=1`
 
 ### Current Models
 - `umx`

cfg/config.yaml

@@ -6,8 +6,8 @@ defaults:
 seed: 12345
 train: True
 sample_rate: 48000
+chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-log_every_n_steps: 1000
 render_files: True
 render_root: "./data/processed"
 
@@ -21,6 +21,9 @@ callbacks:
     verbose: False
     dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
     filename: '{epoch:02d}-{valid_loss:.3f}'
+  learning_rate_monitor:
+    _target_: pytorch_lightning.callbacks.LearningRateMonitor
+    logging_interval: "step"
 
 datamodule:
   _target_: remfx.datasets.VocalSetDatamodule
@@ -28,27 +31,27 @@ datamodule:
     _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
-    chunk_size_in_sec: 6
+    chunk_size: ${chunk_size}
     mode: "train"
-    effect_types: ${effects.train_effects}
+    effect_types: ${effects}
     render_files: ${render_files}
     render_root: ${render_root}
   val_dataset:
     _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
-    chunk_size_in_sec: 6
+    chunk_size: ${chunk_size}
     mode: "val"
-    effect_types: ${effects.val_effects}
+    effect_types: ${effects}
     render_files: ${render_files}
     render_root: ${render_root}
   test_dataset:
     _target_: remfx.datasets.VocalSet
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
-    chunk_size_in_sec: 6
+    chunk_size: ${chunk_size}
     mode: "test"
-    effect_types: ${effects.val_effects}
+    effect_types: ${effects}
     render_files: ${render_files}
     render_root: ${render_root}
 
@@ -76,3 +79,5 @@ trainer:
   accumulate_grad_batches: 1
   accelerator: null
   devices: 1
+  gradient_clip_val: 10.0
+  max_steps: 50000

cfg/effects/all.yaml

@@ -1,70 +1,31 @@
 # @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
+  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/effects/chorus.yaml

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

cfg/effects/compression.yaml

@@ -1,22 +0,0 @@
-# @package _global_
-effects:
-  train_effects:
-      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
-  val_effects:
-      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

cfg/effects/compressor.yaml

@@ -0,0 +1,9 @@
+# @package _global_
+effects:
+  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/effects/distortion.yaml

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

cfg/effects/reverb.yaml

@@ -1,26 +1,13 @@
 # @package _global_
 effects:
-  train_effects:
-      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:
-      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
+  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_compression.yaml → demucs_compressor.yaml}

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

cfg/exp/{umx_compression.yaml → umx_compressor.yaml}

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

remfx/datasets.py

@@ -17,7 +17,7 @@ class VocalSet(Dataset):
         self,
         root: str,
         sample_rate: int,
-        chunk_size_in_sec: int = 3,
+        chunk_size: int = 3,
         effect_types: List[torch.nn.Module] = None,
         render_files: bool = True,
         render_root: str = None,
@@ -28,7 +28,7 @@ class VocalSet(Dataset):
         self.song_idx = []
         self.root = Path(root)
         self.render_root = Path(render_root)
-        self.chunk_size_in_sec = chunk_size_in_sec
+        self.chunk_size = chunk_size
         self.sample_rate = sample_rate
         self.mode = mode
 
@@ -36,9 +36,11 @@ class VocalSet(Dataset):
         self.files = sorted(list(mode_path.glob("./**/*.wav")))
         self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
         self.effect_types = effect_types
-
-        self.processed_root = self.render_root / "processed" / self.mode
-
+        effect_str = "_".join([e for e in self.effect_types])
+        self.processed_root = self.render_root / "processed" / effect_str / self.mode
+        if self.processed_root.exists():
+            print("Found processed files.")
+            render_files = False
         self.num_chunks = 0
         print("Total files:", len(self.files))
         print("Processing files...")
@@ -46,19 +48,14 @@ class VocalSet(Dataset):
             # 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
-                )
+                chunks, orig_sr = create_sequential_chunks(audio_file, self.chunk_size)
                 for chunk in chunks:
                     resampled_chunk = torchaudio.functional.resample(
                         chunk, orig_sr, sample_rate
                     )
-                    chunk_size_in_samples = self.chunk_size_in_sec * self.sample_rate
-                    if resampled_chunk.shape[-1] < chunk_size_in_samples:
-                        resampled_chunk = F.pad(
-                            resampled_chunk,
-                            (0, chunk_size_in_samples - resampled_chunk.shape[1]),
-                        )
+                    if resampled_chunk.shape[-1] < chunk_size:
+                        # Skip if chunk is too small
+                        continue
                     # Apply effect
                     effect_idx = torch.rand(1).item() * len(self.effect_types.keys())
                     effect_name = list(self.effect_types.keys())[int(effect_idx)]

remfx/models.py

@@ -55,6 +55,29 @@ class RemFXModel(pl.LightningModule):
         )
         return optimizer
 
+    # Add step-based learning rate scheduler
+    def optimizer_step(
+        self,
+        epoch,
+        batch_idx,
+        optimizer,
+        optimizer_idx,
+        optimizer_closure,
+        on_tpu,
+        using_native_amp,
+        using_lbfgs,
+    ):
+        # update params
+        optimizer.step(closure=optimizer_closure)
+
+        # update learning rate. Reduce by factor of 10 at 80% and 95% of training
+        if self.trainer.global_step == 0.8 * self.trainer.max_steps:
+            for pg in optimizer.param_groups:
+                pg["lr"] = 0.1 * pg["lr"]
+        if self.trainer.global_step == 0.95 * self.trainer.max_steps:
+            for pg in optimizer.param_groups:
+                pg["lr"] = 0.1 * pg["lr"]
+
     def training_step(self, batch, batch_idx):
         loss = self.common_step(batch, batch_idx, mode="train")
         return loss
@@ -215,7 +238,7 @@ class OpenUnmixModel(torch.nn.Module):
         X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
         Y = self.model(X)
         sep_out = self.separator(x).squeeze(1)
-        loss = self.mrstftloss(sep_out, target) + self.l1loss(sep_out, target)
+        loss = self.mrstftloss(sep_out, target) + self.l1loss(sep_out, target) * 100
 
         return loss, sep_out
 
@@ -236,7 +259,7 @@ class DemucsModel(torch.nn.Module):
     def forward(self, batch):
         x, target, label = batch
         output = self.model(x).squeeze(1)
-        loss = self.mrstftloss(output, target) + self.l1loss(output, target)
+        loss = self.mrstftloss(output, target) + self.l1loss(output, target) * 100
         return loss, output
 
     def sample(self, x: Tensor) -> Tensor:
@@ -264,10 +287,13 @@ def log_wandb_audio_batch(
     samples: Tensor,
     sampling_rate: int,
     caption: str = "",
+    max_items: int = 10,
 ):
     num_items = samples.shape[0]
     samples = rearrange(samples, "b c t -> b t c")
     for idx in range(num_items):
+        if idx >= max_items:
+            break
         logger.experiment.log(
             {
                 f"{id}_{idx}": wandb.Audio(

remfx/utils.py

@@ -132,10 +132,9 @@ def create_sequential_chunks(
     """
     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)
+    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size)
     for start in chunk_starts:
-        if start + chunk_size_in_samples > audio.shape[-1]:
+        if start + chunk_size > audio.shape[-1]:
             break
-        chunks.append(audio[:, start : start + chunk_size_in_samples])
+        chunks.append(audio[:, start : start + chunk_size])
     return chunks, sr

scripts/test.py

@@ -14,7 +14,6 @@ def main(cfg: DictConfig):
     # Apply seed for reproducibility
     if cfg.seed:
         pl.seed_everything(cfg.seed)
-    cfg.render_files = False
     log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>.")
     datamodule = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
     log.info(f"Instantiating model <{cfg.model._target_}>.")

scripts/train.py

@@ -42,6 +42,7 @@ def main(cfg: DictConfig):
     summary = ModelSummary(model)
     print(summary)
     trainer.fit(model=model, datamodule=datamodule)
+    trainer.test(model=model, datamodule=datamodule, ckpt_path="best")
 
 
 if __name__ == "__main__":