Add classifier into inference. Comment-out hearbaseline for now

cfg/exp/chain_inference_aug_classifier.yaml

@@ -0,0 +1,84 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+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
+
+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
+
+classifier:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.Cnn14
+    num_classes: ${num_classes}
+    n_fft: 1024
+    hop_length: 256
+    n_mels: 128
+    sample_rate: 44100
+    model_sample_rate: 44100
+    specaugment: False
+classifier_ckpt: "ckpts/classifier.ckpt"
+
+ckpts:
+  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"
+  - "RandomPedalboardReverb"
+  - "RandomPedalboardChorus"
+  - "RandomPedalboardDelay"
+num_bins: 1025

remfx/classifier.py

@@ -1,9 +1,11 @@
 import torch
 import torchaudio
 import torch.nn as nn
-import hearbaseline
-import hearbaseline.vggish
-import hearbaseline.wav2vec2
+
+# import hearbaseline
+
+# import hearbaseline.vggish
+# import hearbaseline.wav2vec2
 
 import wav2clip_hear
 import panns_hear

remfx/models.py

@@ -20,7 +20,7 @@ ALL_EFFECTS = effects.Pedalboard_Effects
 
 
 class RemFXChainInference(pl.LightningModule):
-    def __init__(self, models, sample_rate, num_bins, effect_order):
+    def __init__(self, models, sample_rate, num_bins, effect_order, classifier=None):
         super().__init__()
         self.model = models
         self.mrstftloss = MultiResolutionSTFTLoss(
@@ -35,6 +35,7 @@ class RemFXChainInference(pl.LightningModule):
         )
         self.sample_rate = sample_rate
         self.effect_order = effect_order
+        self.classifier = classifier
 
     def forward(self, batch, batch_idx, order=None):
         x, y, _, rem_fx_labels = batch
@@ -43,6 +44,13 @@ class RemFXChainInference(pl.LightningModule):
             effects_order = order
         else:
             effects_order = self.effect_order
+
+        # Use classifier labels
+        if self.classifier:
+            threshold = 0.5
+            labels = self.classifier(x)
+            rem_fx_labels = torch.where(labels > threshold, 1.0, 0.0)
+
         effects_present = [
             [ALL_EFFECTS[i] for i, effect in enumerate(effect_label) if effect == 1.0]
             for effect_label in rem_fx_labels

scripts/chain_inference.py

@@ -26,6 +26,16 @@ def main(cfg: DictConfig):
         model.to(device)
         models[effect] = model
 
+    classifier = None
+    if "classifier" in cfg:
+        log.info(f"Instantiating classifier <{cfg.classifier._target_}>.")
+        classifier = hydra.utils.instantiate(cfg.classifier, _convert_="partial")
+        ckpt_path = cfg.classifier_ckpt
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
+        classifier.load_state_dict(state_dict)
+        classifier.to(device)
+
     callbacks = []
     if "callbacks" in cfg:
         for _, cb_conf in cfg["callbacks"].items():
@@ -54,6 +64,7 @@ def main(cfg: DictConfig):
         sample_rate=cfg.sample_rate,
         num_bins=cfg.num_bins,
         effect_order=cfg.inference_effects_ordering,
+        classifier=classifier,
     )
     trainer.test(model=inference_model, datamodule=datamodule)