Add custom-inference dataset

cfg/exp/chain_inference_custom.yaml

@@ -0,0 +1,41 @@
+# @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: 16
+  num_workers: 8
+  train_dataset: None
+  val_dataset: None
+  test_dataset:
+    _target_: remfx.datasets.InferenceDataset
+    root: "./data/fx-examples"
+    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"
+num_bins: 1025

remfx/datasets.py

@@ -360,6 +360,42 @@ 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):
+        self.root = Path(root)
+        self.sample_rate = sample_rate
+        self.clean_paths = list(self.root.glob("clean/*.wav"))
+        self.effected_paths = list(self.root.glob("effected/*.wav"))
+
+    def __len__(self) -> int:
+        return len(self.audio_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)
+        effected = torch.sum(effected, dim=0)
+
+        # Pad or trim effected to clean
+        if len(clean) > len(effected):
+            effected = torch.nn.functional.pad(
+                effected, (0, len(clean) - len(effected))
+            )
+        elif len(effected) > len(clean):
+            effected = effected[: len(clean)]
+
+        dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+        wet_labels_tensor = torch.ones(len(ALL_EFFECTS))
+
+        return clean, effected, dry_labels_tensor, wet_labels_tensor
+
+
 class EffectDatamodule(pl.LightningDataModule):
     def __init__(
         self,

remfx/models.py

@@ -96,7 +96,8 @@ 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 +308,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):