Add effects chain inference code

cfg/config.yaml

@@ -51,7 +51,7 @@ datamodule:
   _target_: remfx.datasets.EffectDatamodule
   train_dataset:
     _target_: remfx.datasets.EffectDataset
-    total_chunks: 8000
+    total_chunks: 80
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -67,7 +67,7 @@ datamodule:
     render_root: ${render_root}
   val_dataset:
     _target_: remfx.datasets.EffectDataset
-    total_chunks: 1000
+    total_chunks: 10
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -83,7 +83,7 @@ datamodule:
     render_root: ${render_root}
   test_dataset:
     _target_: remfx.datasets.EffectDataset
-    total_chunks: 1000
+    total_chunks: 10
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -124,3 +124,10 @@ trainer:
   devices: 1
   gradient_clip_val: 10.0
   max_steps: 50000
+ckpts:
+  RandomPedalboardChorus: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardDelay: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardDistortion: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardCompressor: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardReverb: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+num_bins: 1025

cfg/exp/chain_inference.yaml

@@ -0,0 +1,35 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+render_files: True
+render_root: "/scratch/EffectSet"
+accelerator: "gpu"
+log_audio: True
+# Effects
+num_kept_effects: [0,0] # [min, max]
+num_removed_effects: [0,5] # [min, max]
+shuffle_kept_effects: True
+shuffle_removed_effects: True
+num_classes: 5
+effects_to_keep:
+effects_to_remove:
+  - distortion
+  - compressor
+  - reverb
+  - chorus
+  - delay
+datamodule:
+  batch_size: 16
+  num_workers: 8
+ckpts:
+  RandomPedalboardChorus: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardDelay: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardDistortion: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardCompressor: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+  RandomPedalboardReverb: "/Users/matthewrice/Developer/remfx/ckpts/demucs_5-5/last.ckpt"
+num_bins: 1025

remfx/callbacks.py

@@ -4,6 +4,7 @@ from einops import rearrange
 import torch
 import wandb
 from torch import Tensor
+from remfx.models import RemFXChainInference
 
 
 class AudioCallback(Callback):
@@ -46,7 +47,10 @@ class AudioCallback(Callback):
         # Only run on first batch
         if batch_idx == 0 and self.log_audio:
             with torch.no_grad():
-                y = pl_module.model.sample(x)
+                if type(pl_module) == RemFXChainInference:
+                    y = pl_module.sample(batch)
+                else:
+                    y = pl_module.model.sample(x)
             # Concat samples together for easier viewing in dashboard
             # 2 seconds of silence between each sample
             silence = torch.zeros_like(x)

remfx/models.py

@@ -11,8 +11,78 @@ from umx.openunmix.model import OpenUnmix, Separator
 from remfx.utils import FADLoss, spectrogram
 from remfx.tcn import TCN
 from remfx.utils import causal_crop
+from remfx import effects
 import asteroid
 
+ALL_EFFECTS = effects.Pedalboard_Effects
+
+
+class RemFXChainInference(pl.LightningModule):
+    def __init__(self, models, sample_rate, num_bins):
+        super().__init__()
+        self.model = models
+        self.mrstftloss = MultiResolutionSTFTLoss(
+            n_bins=num_bins, sample_rate=sample_rate
+        )
+        self.l1loss = nn.L1Loss()
+        self.metrics = nn.ModuleDict(
+            {
+                "SISDR": SISDRLoss(),
+                "STFT": MultiResolutionSTFTLoss(),
+                "FAD": FADLoss(sample_rate=sample_rate),
+            }
+        )
+
+    def forward(self, batch):
+        x, y, _, rem_fx_labels = batch
+        # Use chain of effects defined in config
+        effects = [
+            [ALL_EFFECTS[i] for i, effect in enumerate(effect_label) if effect == 1.0]
+            for effect_label in rem_fx_labels
+        ]
+        output = []
+        with torch.no_grad():
+            for elem, effect_chain in zip(x, effects):
+                elem = elem.unsqueeze(0)  # Add batch dim
+                for effect in effect_chain:
+                    # Get correct model based on effect name. This is a bit hacky
+                    # Then sample the model
+                    elem = self.model[effect.__name__].model.sample(elem)
+                output.append(elem.squeeze(0))
+        output = torch.stack(output)
+
+        loss = self.mrstftloss(output, y) + self.l1loss(output, y) * 100
+        return loss, output
+
+    def test_step(self, batch, batch_idx):
+        x, y, _, _ = batch  # x, y = (B, C, T), (B, C, T)
+
+        loss, output = self.forward(batch)
+        # Crop target to match output
+        if output.shape[-1] < y.shape[-1]:
+            y = causal_crop(y, output.shape[-1])
+        self.log("test_loss", loss)
+        # Metric logging
+        with torch.no_grad():
+            for metric in self.metrics:
+                # SISDR returns negative values, so negate them
+                if metric == "SISDR":
+                    negate = -1
+                else:
+                    negate = 1
+                self.log(
+                    f"test_{metric}",
+                    negate * self.metrics[metric](output, y),
+                    on_step=False,
+                    on_epoch=True,
+                    logger=True,
+                    prog_bar=True,
+                    sync_dist=True,
+                )
+
+    def sample(self, batch):
+        return self.forward(batch)[1]
+
 
 class RemFX(pl.LightningModule):
     def __init__(

scripts/chain_inference.py

@@ -0,0 +1,61 @@
+import pytorch_lightning as pl
+import hydra
+from omegaconf import DictConfig
+import remfx.utils as utils
+from pytorch_lightning.utilities.model_summary import ModelSummary
+import torch
+from remfx.models import RemFXChainInference
+
+log = utils.get_logger(__name__)
+
+
+@hydra.main(version_base=None, config_path="../cfg", config_name="config.yaml")
+def main(cfg: DictConfig):
+    # Apply seed for reproducibility
+    if cfg.seed:
+        pl.seed_everything(cfg.seed)
+    log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>.")
+    datamodule = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
+    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")
+        state_dict = torch.load(ckpt_path, map_location=torch.device("cpu"))[
+            "state_dict"
+        ]
+        model.load_state_dict(state_dict)
+        models[effect] = model
+
+    callbacks = []
+    if "callbacks" in cfg:
+        for _, cb_conf in cfg["callbacks"].items():
+            if "_target_" in cb_conf:
+                log.info(f"Instantiating callback <{cb_conf._target_}>.")
+                callbacks.append(hydra.utils.instantiate(cb_conf, _convert_="partial"))
+
+    logger = hydra.utils.instantiate(cfg.logger, _convert_="partial")
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>.")
+    trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=logger, _convert_="partial"
+    )
+    log.info("Logging hyperparameters!")
+    utils.log_hyperparameters(
+        config=cfg,
+        model=model,
+        datamodule=datamodule,
+        trainer=trainer,
+        callbacks=callbacks,
+        logger=logger,
+    )
+    summary = ModelSummary(model)
+    print(summary)
+
+    inference_model = RemFXChainInference(
+        models, sample_rate=cfg.sample_rate, num_bins=cfg.num_bins
+    )
+    trainer.test(model=inference_model, datamodule=datamodule)
+
+
+if __name__ == "__main__":
+    main()