Merge pull request #38 from mhrice/custom-inference

README.md

@@ -47,6 +47,23 @@ see `cfg/exp/default.yaml` for an example.
 - `reverb`
 - `delay`
 
+## Run inference on directory
+Assumes directory is structured as
+- root
+    - clean
+        - file1.wav
+        - file2.wav
+        - file3.wav
+    - effected
+        - file1.wav
+        - file2.wav
+        - file3.wav
+
+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

@@ -28,10 +28,17 @@ datamodule:
   num_workers: 8
   train_dataset: None
   val_dataset: None
+
 ckpts:
   RandomPedalboardDistortion: "ckpts/distortion.ckpt"
   RandomPedalboardCompressor: "ckpts/compressor.ckpt"
   RandomPedalboardReverb: "ckpts/reverb.ckpt"
   RandomPedalboardChorus: "ckpts/chorus.ckpt"
   RandomPedalboardDelay: "ckpts/delay.ckpt"
+inference_effects_ordering:
+  - "RandomPedalboardDistortion"
+  - "RandomPedalboardCompressor"
+  - "RandomPedalboardReverb"
+  - "RandomPedalboardChorus"
+  - "RandomPedalboardDelay"
 num_bins: 1025

cfg/exp/chain_inference_custom.yaml

@@ -0,0 +1,47 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+render_files: False
+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: 1
+  num_workers: 8
+  train_dataset: None
+  val_dataset: None
+  test_dataset:
+    _target_: remfx.datasets.InferenceDataset
+    root: ${oc.env:DATASET_ROOT}
+    sample_rate: ${sample_rate}
+ckpts:
+  RandomPedalboardDistortion: "ckpts/distortion.ckpt"
+  RandomPedalboardCompressor: "ckpts/compressor.ckpt"
+  RandomPedalboardReverb: "ckpts/reverb.ckpt"
+  RandomPedalboardChorus: "ckpts/chorus.ckpt"
+  RandomPedalboardDelay: "ckpts/delay.ckpt"
+inference_effects_ordering:
+  - "RandomPedalboardDistortion"
+  - "RandomPedalboardCompressor"
+  - "RandomPedalboardReverb"
+  - "RandomPedalboardChorus"
+  - "RandomPedalboardDelay"
+num_bins: 1025

remfx/datasets.py

@@ -360,6 +360,41 @@ class EffectDataset(Dataset):
         return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
 
 
+class InferenceDataset(Dataset):
+    def __init__(self, root: str, sample_rate: int, **kwargs):
+        self.root = Path(root)
+        self.sample_rate = sample_rate
+        self.clean_paths = sorted(list(self.root.glob("clean/*.wav")))
+        self.effected_paths = sorted(list(self.root.glob("effected/*.wav")))
+
+    def __len__(self) -> int:
+        return len(self.clean_paths)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        clean_path = self.clean_paths[idx]
+        effected_path = self.effected_paths[idx]
+        clean_audio, sr = torchaudio.load(clean_path)
+        clean = torchaudio.functional.resample(clean_audio, sr, self.sample_rate)
+        effected_audio, sr = torchaudio.load(effected_path)
+        effected = torchaudio.functional.resample(effected_audio, sr, self.sample_rate)
+
+        # Sum to mono
+        clean = torch.sum(clean, dim=0, keepdim=True)
+        effected = torch.sum(effected, dim=0, keepdim=True)
+
+        # Pad or trim effected to clean
+        if effected.shape[1] > clean.shape[1]:
+            effected = effected[:, : clean.shape[1]]
+        elif effected.shape[1] < clean.shape[1]:
+            pad_size = clean.shape[1] - effected.shape[1]
+            effected = torch.nn.functional.pad(effected, (0, pad_size))
+
+        dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+        wet_labels_tensor = torch.ones(len(ALL_EFFECTS))
+
+        return effected, clean, dry_labels_tensor, wet_labels_tensor
+
+
 class EffectDatamodule(pl.LightningDataModule):
     def __init__(
         self,

remfx/models.py

@@ -37,7 +37,7 @@ class RemFXChainInference(pl.LightningModule):
         self.sample_rate = sample_rate
         self.effect_order = effect_order
 
-    def forward(self, batch, order=None):
+    def forward(self, batch, batch_idx, order=None):
         x, y, _, rem_fx_labels = batch
         # Use chain of effects defined in config
         if order:
@@ -52,25 +52,30 @@ class RemFXChainInference(pl.LightningModule):
         with torch.no_grad():
             for i, (elem, effects_list) in enumerate(zip(x, effects)):
                 elem = elem.unsqueeze(0)  # Add batch dim
-                effect_chain_idx = [
-                    effects_order.index(effect.__name__) for effect in effects_list
+                # effect_chain_idx = [
+                #     effects_order.index(effect.__name__) for effect in effects_list
+                # ]
+                effect_list_names = [effect.__name__ for effect in effects_list]
+                effects = [
+                    effect for effect in effects_order if effect in effect_list_names
                 ]
+
                 # log_wandb_audio_batch(
                 #     logger=self.logger,
-                #     id=f"{i}_Before",
+                #     id=f"{batch_idx}_{i}_Before",
                 #     samples=elem.cpu(),
                 #     sampling_rate=self.sample_rate,
-                #     caption=effect_chain,
+                #     caption=effects,
                 # )
-                for idx in effect_chain_idx:
+                for effect in effects:
                     # Sample the model
-                    elem = self.model[effects_order[idx]].model.sample(elem)
+                    elem = self.model[effect].model.sample(elem)
                     # log_wandb_audio_batch(
                     #     logger=self.logger,
-                    #     id=f"{i}_{effect}",
+                    #     id=f"{batch_idx}_{i}_{effect}",
                     #     samples=elem.cpu(),
                     #     sampling_rate=self.sample_rate,
-                    #     caption=effect_chain,
+                    #     caption=effects,
                     # )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
@@ -81,8 +86,8 @@ class RemFXChainInference(pl.LightningModule):
     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)
-        loss, output = self.forward(batch, order=self.effect_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]:
             y = causal_crop(y, output.shape[-1])
@@ -96,7 +101,7 @@ class RemFXChainInference(pl.LightningModule):
                 else:
                     negate = 1
                 self.log(
-                    f"test_{metric}_" + "".join(self.effect_order),
+                    f"test_{metric}",  # + "".join(self.effect_order).replace("RandomPedalboard", ""),
                     negate * self.metrics[metric](output, y),
                     on_step=False,
                     on_epoch=True,
@@ -307,27 +312,6 @@ class DPTNetModel(nn.Module):
     def sample(self, x: Tensor) -> Tensor:
         return self.model(x.squeeze(1))
 
-    def __init__(self, sample_rate, num_bins, **kwargs):
-        super().__init__()
-        self.model = asteroid.models.DCUNet(**kwargs)
-        self.mrstftloss = MultiResolutionSTFTLoss(
-            n_bins=num_bins, sample_rate=sample_rate
-        )
-        self.l1loss = nn.L1Loss()
-
-    def forward(self, batch):
-        x, target = batch
-        output = self.model(x.squeeze(1))  # B x T
-        # Crop target to match output
-        if output.shape[-1] < target.shape[-1]:
-            target = causal_crop(target, output.shape[-1])
-        loss = self.mrstftloss(output, target) + self.l1loss(output, target) * 100
-        return loss, output
-
-    def sample(self, x: Tensor) -> Tensor:
-        output = self.model(x.squeeze(1))  # B x T
-        return output
-
 
 class TCNModel(nn.Module):
     def __init__(self, sample_rate, num_bins, **kwargs):

scripts/chain_inference.py

@@ -47,17 +47,12 @@ def main(cfg: DictConfig):
         logger=logger,
     )
 
+    log.info("Instantiating Inference Model")
     inference_model = RemFXChainInference(
         models,
         sample_rate=cfg.sample_rate,
         num_bins=cfg.num_bins,
-        effect_order=[
-            "RandomPedalboardDistortion",
-            "RandomPedalboardCompressor",
-            "RandomPedalboardReverb",
-            "RandomPedalboardChorus",
-            "RandomPedalboardDelay",
-        ],
+        effect_order=cfg.inference_effects_ordering,
     )
     trainer.test(model=inference_model, datamodule=datamodule)