Merge pull request #40 from mhrice/classifier-inference

README.md

@@ -47,6 +47,9 @@ see `cfg/exp/default.yaml` for an example.
 - `reverb`
 - `delay`
 
+## Chain Inference
+`python scripts/chain_inference.py +exp=chain_inference`
+
 ## Run inference on directory
 Assumes directory is structured as
 - root
@@ -64,6 +67,7 @@ Change root path in `shell_vars.sh` and `source shell_vars.sh`
 `python scripts/chain_inference.py +exp=chain_inference_custom`
 
 
+
 ## Misc.
 By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
 

cfg/exp/chain_inference.yaml

@@ -63,4 +63,6 @@ inference_effects_ordering:
   - "RandomPedalboardReverb"
   - "RandomPedalboardChorus"
   - "RandomPedalboardDelay"
-num_bins: 1025
+num_bins: 1025
+inference_effects_shuffle: False
+inference_use_all_effect_models: False

cfg/exp/chain_inference_aug.yaml

@@ -63,4 +63,6 @@ inference_effects_ordering:
   - "RandomPedalboardReverb"
   - "RandomPedalboardChorus"
   - "RandomPedalboardDelay"
-num_bins: 1025
+num_bins: 1025
+inference_effects_shuffle: False
+inference_use_all_effect_models: False

cfg/exp/chain_inference_aug_classifier.yaml

@@ -0,0 +1,87 @@
+# @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}
+  mixup: False
+  network:
+    _target_: remfx.classifier.Cnn14
+    num_classes: ${num_classes}
+    n_fft: 2048
+    hop_length: 512
+    n_mels: 128
+    sample_rate: ${sample_rate}
+    model_sample_rate: ${sample_rate}
+    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
+inference_effects_shuffle: False
+inference_use_all_effect_models: False

cfg/exp/chain_inference_custom.yaml

@@ -68,4 +68,6 @@ inference_effects_ordering:
   - "RandomPedalboardReverb"
   - "RandomPedalboardChorus"
   - "RandomPedalboardDelay"
-num_bins: 1025
+num_bins: 1025
+inference_effects_shuffle: False
+inference_use_all_effect_models: False

remfx/callbacks.py

@@ -64,9 +64,6 @@ class AudioCallback(Callback):
                     ]
                     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

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/datasets.py

@@ -13,10 +13,10 @@ from typing import Any, List, Dict
 from torch.utils.data import Dataset, DataLoader
 from remfx.utils import select_random_chunk
 import multiprocessing
+from auraloss.freq import MultiResolutionSTFTLoss
 
 
-# https://zenodo.org/record/1193957 -> VocalSet
-
+STFT_THRESH = 1e-3
 ALL_EFFECTS = effect_lib.Pedalboard_Effects
 # print(ALL_EFFECTS)
 
@@ -404,6 +404,7 @@ class EffectDataset(Dataset):
         self.effects_to_keep = [] if effects_to_keep is None else effects_to_keep
         self.effects_to_remove = [] if effects_to_remove is None else effects_to_remove
         self.normalize = effect_lib.LoudnessNormalize(sample_rate, target_lufs_db=-20)
+        self.mrstft = MultiResolutionSTFTLoss(sample_rate=sample_rate)
         self.effects = effect_modules
         self.shuffle_kept_effects = shuffle_kept_effects
         self.shuffle_removed_effects = shuffle_removed_effects
@@ -471,7 +472,6 @@ class EffectDataset(Dataset):
                         chunk = select_random_chunk(
                             random_file_choice, self.chunk_size, self.sample_rate
                         )
-
                     # Sum to mono
                     if chunk.shape[0] > 1:
                         chunk = chunk.sum(0, keepdim=True)
@@ -568,46 +568,52 @@ class EffectDataset(Dataset):
         # 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]
-        # Apply
-        dry_labels = []
-        for effect in effects_to_apply:
-            # Normalize in-between effects
-            dry = self.normalize(effect(dry))
-            dry_labels.append(ALL_EFFECTS.index(type(effect)))
-
-        # Apply effects_to_remove
-        # Shuffle effects if specified
-        if self.shuffle_removed_effects:
-            effect_indices = torch.randperm(len(self.effects_to_remove))
-        else:
-            effect_indices = torch.arange(len(self.effects_to_remove))
-        wet = torch.clone(dry)
-        r1 = self.num_removed_effects[0]
-        r2 = self.num_removed_effects[1]
-        num_removed_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
-        effect_indices = effect_indices[:num_removed_effects]
-        # Index in effect settings
-        effect_names_to_apply = [self.effects_to_remove[i] for i in effect_indices]
-        effects_to_apply = [self.effects[i] for i in effect_names_to_apply]
-        # Apply
-        wet_labels = []
-        for effect in effects_to_apply:
-            # Normalize in-between effects
-            wet = self.normalize(effect(wet))
-            wet_labels.append(ALL_EFFECTS.index(type(effect)))
-
-        wet_labels_tensor = torch.zeros(len(ALL_EFFECTS))
-        dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
-
-        for label_idx in wet_labels:
-            wet_labels_tensor[label_idx] = 1.0
-
-        for label_idx in dry_labels:
-            dry_labels_tensor[label_idx] = 1.0
-
-        # Normalize
-        normalized_dry = self.normalize(dry)
-        normalized_wet = self.normalize(wet)
+        # stft comparison
+        stft = 0
+        while stft < STFT_THRESH:
+            # Apply
+            dry_labels = []
+            for effect in effects_to_apply:
+                # Normalize in-between effects
+                dry = self.normalize(effect(dry))
+                dry_labels.append(ALL_EFFECTS.index(type(effect)))
+
+            # Apply effects_to_remove
+            # Shuffle effects if specified
+            if self.shuffle_removed_effects:
+                effect_indices = torch.randperm(len(self.effects_to_remove))
+            else:
+                effect_indices = torch.arange(len(self.effects_to_remove))
+            wet = torch.clone(dry)
+            r1 = self.num_removed_effects[0]
+            r2 = self.num_removed_effects[1]
+            num_removed_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
+            effect_indices = effect_indices[:num_removed_effects]
+            # Index in effect settings
+            effect_names_to_apply = [self.effects_to_remove[i] for i in effect_indices]
+            effects_to_apply = [self.effects[i] for i in effect_names_to_apply]
+            # Apply
+            wet_labels = []
+            for effect in effects_to_apply:
+                # Normalize in-between effects
+                wet = self.normalize(effect(wet))
+                wet_labels.append(ALL_EFFECTS.index(type(effect)))
+
+            wet_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+            dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+
+            for label_idx in wet_labels:
+                wet_labels_tensor[label_idx] = 1.0
+
+            for label_idx in dry_labels:
+                dry_labels_tensor[label_idx] = 1.0
+
+            # Normalize
+            normalized_dry = self.normalize(dry)
+            normalized_wet = self.normalize(wet)
+
+            # Check STFT, pick different effects if necessary
+            stft = self.mrstft(normalized_wet, normalized_dry)
         return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
 
 
@@ -692,7 +698,7 @@ class EffectDatamodule(pl.LightningDataModule):
     def test_dataloader(self) -> DataLoader:
         return DataLoader(
             dataset=self.test_dataset,
-            batch_size=self.test_batch_size,  # Use small, consistent batch size for testing
+            batch_size=1,  # Use small, consistent batch size for testing
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             shuffle=False,

remfx/models.py

@@ -16,12 +16,22 @@ from remfx.callbacks import log_wandb_audio_batch
 from einops import rearrange
 from remfx import effects
 import asteroid
+import random
 
 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,
+        shuffle_effect_order=False,
+        use_all_effect_models=False,
+    ):
         super().__init__()
         self.model = models
         self.mrstftloss = MultiResolutionSTFTLoss(
@@ -36,6 +46,10 @@ class RemFXChainInference(pl.LightningModule):
         )
         self.sample_rate = sample_rate
         self.effect_order = effect_order
+        self.classifier = classifier
+        self.shuffle_effect_order = shuffle_effect_order
+        self.output_str = "IN_SISDR,OUT_SISDR,IN_STFT,OUT_STFT\n"
+        self.use_all_effect_models = use_all_effect_models
 
     def forward(self, batch, batch_idx, order=None):
         x, y, _, rem_fx_labels = batch
@@ -44,28 +58,46 @@ class RemFXChainInference(pl.LightningModule):
             effects_order = order
         else:
             effects_order = self.effect_order
-        effects_present = [
-            [ALL_EFFECTS[i] for i, effect in enumerate(effect_label) if effect == 1.0]
-            for effect_label in rem_fx_labels
-        ]
+
+        # Use classifier labels
+        if self.classifier:
+            threshold = 0.5
+            with torch.no_grad():
+                labels = torch.sigmoid(self.classifier(x))
+                rem_fx_labels = torch.where(labels > threshold, 1.0, 0.0)
+        if self.use_all_effect_models:
+            effects_present = [
+                [ALL_EFFECTS[i] for i, effect in enumerate(effect_label)]
+                for effect_label in rem_fx_labels
+            ]
+        else:
+            effects_present = [
+                [
+                    ALL_EFFECTS[i]
+                    for i, effect in enumerate(effect_label)
+                    if effect == 1.0
+                ]
+                for effect_label in rem_fx_labels
+            ]
+
         output = []
-        input_samples = rearrange(x, "b c t -> c (b t)").unsqueeze(0)
-        target_samples = rearrange(y, "b c t -> c (b t)").unsqueeze(0)
-
-        log_wandb_audio_batch(
-            logger=self.logger,
-            id="input_effected_audio",
-            samples=input_samples.cpu(),
-            sampling_rate=self.sample_rate,
-            caption="Input Data",
-        )
-        log_wandb_audio_batch(
-            logger=self.logger,
-            id="target_audio",
-            samples=target_samples.cpu(),
-            sampling_rate=self.sample_rate,
-            caption="Target Data",
-        )
+        # input_samples = rearrange(x, "b c t -> c (b t)").unsqueeze(0)
+        # target_samples = rearrange(y, "b c t -> c (b t)").unsqueeze(0)
+
+        # log_wandb_audio_batch(
+        #     logger=self.logger,
+        #     id="input_effected_audio",
+        #     samples=input_samples.cpu(),
+        #     sampling_rate=self.sample_rate,
+        #     caption="Input Data",
+        # )
+        # log_wandb_audio_batch(
+        #     logger=self.logger,
+        #     id="target_audio",
+        #     samples=target_samples.cpu(),
+        #     sampling_rate=self.sample_rate,
+        #     caption="Target Data",
+        # )
         with torch.no_grad():
             for i, (elem, effects_list) in enumerate(zip(x, effects_present)):
                 elem = elem.unsqueeze(0)  # Add batch dim
@@ -101,22 +133,22 @@ class RemFXChainInference(pl.LightningModule):
                 # )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
-        output_samples = rearrange(output, "b c t -> c (b t)").unsqueeze(0)
-
-        log_wandb_audio_batch(
-            logger=self.logger,
-            id="output_audio",
-            samples=output_samples.cpu(),
-            sampling_rate=self.sample_rate,
-            caption="Output Data",
-        )
+
+        # log_wandb_audio_batch(
+        #     logger=self.logger,
+        #     id="output_audio",
+        #     samples=output_samples.cpu(),
+        #     sampling_rate=self.sample_rate,
+        #     caption="Output Data",
+        # )
         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)
-        # Random order
-        # random.shuffle(self.effect_order)
+        if self.shuffle_effect_order:
+            # Random order
+            random.shuffle(self.effect_order)
         loss, output = self.forward(batch, batch_idx, order=self.effect_order)
         # Crop target to match output
         if output.shape[-1] < y.shape[-1]:
@@ -148,8 +180,16 @@ class RemFXChainInference(pl.LightningModule):
                     prog_bar=True,
                     sync_dist=True,
                 )
+                # print(f"Input_{metric}", negate * self.metrics[metric](x, y))
+                # print(f"test_{metric}", negate * self.metrics[metric](output, y))
+                self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
+            self.output_str += "\n"
         return loss
 
+    def on_test_end(self) -> None:
+        with open("output.csv", "w") as f:
+            f.write(self.output_str)
+
     def sample(self, batch):
         return self.forward(batch, 0)[1]
 
@@ -181,6 +221,7 @@ class RemFX(pl.LightningModule):
         )
         # Log first batch metrics input vs output only once
         self.log_train_audio = True
+        self.output_str = "IN_SISDR,OUT_SISDR,IN_STFT,OUT_STFT\n"
 
     @property
     def device(self):
@@ -257,9 +298,16 @@ class RemFX(pl.LightningModule):
                     prog_bar=True,
                     sync_dist=True,
                 )
-
+                # print(f"Input_{metric}", negate * self.metrics[metric](x, y))
+                # print(f"test_{metric}", negate * self.metrics[metric](output, y))
+                self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
+            self.output_str += "\n"
         return loss
 
+    def on_test_end(self) -> None:
+        with open("output.csv", "w") as f:
+            f.write(self.output_str)
+
 
 class OpenUnmixModel(nn.Module):
     def __init__(

remfx/utils.py

@@ -159,7 +159,7 @@ def select_random_chunk(
     random_start = torch.randint(0, max_len, (1,)).item()
     chunk = audio[:, random_start : random_start + new_chunk_size]
     # Skip if energy too low
-    if torch.mean(torch.abs(chunk)) < 1e-6:
+    if torch.mean(torch.abs(chunk)) < 1e-4:
         return None
     resampled_chunk = torchaudio.functional.resample(chunk, sr, sample_rate)
     return resampled_chunk

scripts/chain_inference.py

@@ -15,7 +15,7 @@ def main(cfg: DictConfig):
         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_}>.")
+    log.info("Instantiating Chain Inference Models")
     models = {}
     for effect in cfg.ckpts:
         model = hydra.utils.instantiate(cfg.ckpts[effect].model, _convert_="partial")
@@ -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,9 @@ def main(cfg: DictConfig):
         sample_rate=cfg.sample_rate,
         num_bins=cfg.num_bins,
         effect_order=cfg.inference_effects_ordering,
+        classifier=classifier,
+        shuffle_effect_order=cfg.inference_effects_shuffle,
+        use_all_effect_models=cfg.inference_use_all_effect_models,
     )
     trainer.test(model=inference_model, datamodule=datamodule)