Add custom model choice for chain inference

cfg/exp/chain_inference.yaml

@@ -26,12 +26,37 @@ datamodule:
   batch_size: 16
   num_workers: 8
 
+dcunet:
+  _target_: remfx.models.RemFX
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.models.DCUNetModel
+    architecture: "Large-DCUNet-20"
+    stft_kernel_size: 512
+    fix_length_mode: "pad"
+    sample_rate: ${sample_rate}
+    num_bins: 1025
 ckpts:
-  RandomPedalboardDistortion: "ckpts/demucs_distortion.ckpt"
-  RandomPedalboardCompressor: "ckpts/demucs_compressor.ckpt"
-  RandomPedalboardReverb: "ckpts/demucs_reverb.ckpt"
-  RandomPedalboardChorus: "ckpts/demucs_chorus.ckpt"
-  RandomPedalboardDelay: "ckpts/demucs_delay.ckpt"
+  RandomPedalboardDistortion:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_distortion.ckpt"
+  RandomPedalboardCompressor:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_compressor.ckpt"
+  RandomPedalboardReverb:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_reverb.ckpt"
+  RandomPedalboardChorus:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_chorus.ckpt"
+  RandomPedalboardDelay:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_delay.ckpt"
 inference_effects_ordering:
   - "RandomPedalboardDistortion"
   - "RandomPedalboardCompressor"

cfg/exp/chain_inference_aug.yaml

@@ -26,12 +26,37 @@ datamodule:
   batch_size: 16
   num_workers: 8
 
+dcunet:
+  _target_: remfx.models.RemFX
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.models.DCUNetModel
+    architecture: "Large-DCUNet-20"
+    stft_kernel_size: 512
+    fix_length_mode: "pad"
+    sample_rate: ${sample_rate}
+    num_bins: 1025
 ckpts:
-  RandomPedalboardDistortion: "ckpts/demucs_distortion_aug.ckpt"
-  RandomPedalboardCompressor: "ckpts/demucs_compressor_aug.ckpt"
-  RandomPedalboardReverb: "ckpts/demucs_reverb_aug.ckpt"
-  RandomPedalboardChorus: "ckpts/demucs_chorus_aug.ckpt"
-  RandomPedalboardDelay: "ckpts/demucs_delay_aug.ckpt"
+  RandomPedalboardDistortion:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_distortion_aug.ckpt"
+  RandomPedalboardCompressor:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_compressor_aug.ckpt"
+  RandomPedalboardReverb:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_reverb_aug.ckpt"
+  RandomPedalboardChorus:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_chorus_aug.ckpt"
+  RandomPedalboardDelay:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_delay_aug.ckpt"
 inference_effects_ordering:
   - "RandomPedalboardDistortion"
   - "RandomPedalboardCompressor"

cfg/exp/chain_inference_custom.yaml

@@ -31,12 +31,37 @@ datamodule:
     _target_: remfx.datasets.InferenceDataset
     root: ${oc.env:DATASET_ROOT}
     sample_rate: ${sample_rate}
+dcunet:
+  _target_: remfx.models.RemFX
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.models.DCUNetModel
+    architecture: "Large-DCUNet-20"
+    stft_kernel_size: 512
+    fix_length_mode: "pad"
+    sample_rate: ${sample_rate}
+    num_bins: 1025
 ckpts:
-  RandomPedalboardDistortion: "ckpts/demucs_distortion_aug.ckpt"
-  RandomPedalboardCompressor: "ckpts/demucs_compressor_aug.ckpt"
-  RandomPedalboardReverb: "ckpts/demucs_reverb_aug.ckpt"
-  RandomPedalboardChorus: "ckpts/demucs_chorus_aug.ckpt"
-  RandomPedalboardDelay: "ckpts/demucs_delay_aug.ckpt"
+  RandomPedalboardDistortion:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_distortion_aug.ckpt"
+  RandomPedalboardCompressor:
+    model: ${model}
+    ckpt_path: "ckpts/demucs_compressor_aug.ckpt"
+  RandomPedalboardReverb:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_reverb_aug.ckpt"
+  RandomPedalboardChorus:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_chorus_aug.ckpt"
+  RandomPedalboardDelay:
+    model: ${dcunet}
+    ckpt_path: "ckpts/dcunet_delay_aug.ckpt"
 inference_effects_ordering:
   - "RandomPedalboardDistortion"
   - "RandomPedalboardCompressor"

remfx/callbacks.py

@@ -4,6 +4,9 @@ from einops import rearrange
 import torch
 import wandb
 from torch import Tensor
+from remfx import effects
+
+ALL_EFFECTS = effects.Pedalboard_Effects
 
 
 class AudioCallback(Callback):
@@ -42,7 +45,7 @@ class AudioCallback(Callback):
     def on_validation_batch_start(
         self, trainer, pl_module, batch, batch_idx, dataloader_idx
     ):
-        x, target, _, _ = batch
+        x, target, _, rem_fx_labels = batch
         # Only run on first batch
         if batch_idx == 0 and self.log_audio:
             with torch.no_grad():
@@ -51,6 +54,19 @@ class AudioCallback(Callback):
 
                 if type(pl_module) == RemFXChainInference:
                     y = pl_module.sample(batch)
+                    effects_present_name = [
+                        [
+                            ALL_EFFECTS[i].__name__.replace("RandomPedalboard", "")
+                            for i, effect in enumerate(effect_label)
+                            if effect == 1.0
+                        ]
+                        for effect_label in rem_fx_labels
+                    ]
+                    for i, label in enumerate(effects_present_name):
+                        self.log(f"{'_'.join(label)}", 0.0)
+                        # self.log(f"{effects}_{i}", label)
+                        # trainer.logger.experiment.log(
+                        # {f"effects_{i}": f"{'_'.join(label)}"}
                 else:
                     y = pl_module.model.sample(x)
             # Concat samples together for easier viewing in dashboard

scripts/chain_inference.py

@@ -18,8 +18,8 @@ def main(cfg: DictConfig):
     log.info(f"Instantiating model <{cfg.model._target_}>.")
     models = {}
     for effect in cfg.ckpts:
-        ckpt_path = cfg.ckpts[effect]
-        model = hydra.utils.instantiate(cfg.model, _convert_="partial")
+        model = hydra.utils.instantiate(cfg.ckpts[effect].model, _convert_="partial")
+        ckpt_path = cfg.ckpts[effect].ckpt_path
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
         model.load_state_dict(state_dict)

scripts/train.py

@@ -18,7 +18,16 @@ def main(cfg: DictConfig):
 
     if "ckpt_path" in cfg:
         log.info(f"Loading checkpoint from <{cfg.ckpt_path}>.")
-        model = model.load_from_checkpoint(cfg.ckpt_path)
+        model.load_from_checkpoint(
+            cfg.ckpt_path,
+            lr=model.lr,
+            lr_beta1=model.lr_beta1,
+            lr_beta2=model.lr_beta2,
+            lr_eps=model.lr_eps,
+            lr_weight_decay=model.lr_weight_decay,
+            sample_rate=model.sample_rate,
+            network=model.model,
+        )
 
     # Init all callbacks
     callbacks = []