Merge pull request #39 from mhrice/cjs--classifier-v2

README.md

@@ -77,4 +77,30 @@ python scripts/download.py vocalset guitarset idmt-smt-guitar idmt-smt-bass idmt
 To run audio effects classifiction:
 ```
 python scripts/train.py model=classifier "effects_to_use=[compressor, distortion, reverb, chorus, delay]" "effects_to_remove=[]" max_kept_effects=5 max_removed_effects=0 shuffle_kept_effects=True shuffle_removed_effects=True accelerator='gpu' render_root=/scratch/RemFX render_files=True
-```
+```
+
+```
+srun --comment harmonai --partition=g40 --gpus=1 --cpus-per-gpu=12 --job-name=harmonai --pty bash -i
+source env/bin/activate
+rsync -aP /fsx/home-csteinmetz1/data/EffectSet_cjs.tar /scratch
+tar -xvf EffectSet_cjs.tar
+mv scratch/EffectSet_cjs ./EffectSet_cjs
+
+export DATASET_ROOT="/admin/home-csteinmetz1/data/remfx-data"
+export WANDB_PROJECT="RemFX"
+export WANDB_ENTITY="cjstein"
+
+python scripts/train.py +exp=5-5.yaml model=cls_vggish render_files=False logs_dir=/scratch/cjs-log datamodule.batch_size=64
+python scripts/train.py +exp=5-5.yaml model=cls_panns_pt render_files=False logs_dir=/scratch/cjs-log datamodule.batch_size=64
+python scripts/train.py +exp=5-5.yaml model=cls_wav2vec2 render_files=False logs_dir=/scratch/cjs-log datamodule.batch_size=64
+python scripts/train.py +exp=5-5.yaml model=cls_wav2clip render_files=False logs_dir=/scratch/cjs-log datamodule.batch_size=64
+```
+
+### Installing HEAR models
+
+wav2clip
+```
+pip install hearbaseline
+pip install git+https://github.com/hohsiangwu/wav2clip-hear.git
+pip install git+https://github.com/qiuqiangkong/HEAR2021_Challenge_PANNs
+wget https://zenodo.org/record/6332525/files/hear2021-panns_hear.pth

cfg/config.yaml

@@ -63,6 +63,7 @@ datamodule:
     shuffle_removed_effects: ${shuffle_removed_effects}
     render_files: ${render_files}
     render_root: ${render_root}
+    parallel: True
   val_dataset:
     _target_: remfx.datasets.EffectDataset
     total_chunks: 1000
@@ -109,6 +110,7 @@ logger:
   job_type: "train"
   group: ""
   save_dir: "."
+  log_model: True
 
 trainer:
   _target_: pytorch_lightning.Trainer

cfg/exp/5-5_cls.yaml

@@ -0,0 +1,60 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "/scratch/cjs-logs"
+render_files: True
+render_root: "/scratch/EffectSet_cjs"
+accelerator: "gpu"
+log_audio: False
+# 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:
+  train_batch_size: 64
+  test_batch_size: 256
+  num_workers: 8
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: "valid_avg_acc_epoch"   # name of the logged metric which determines when model is improving
+    save_top_k: 1           # save k best models (determined by above metric)
+    save_last: True         # additionaly always save model from last epoch
+    mode: "max"             # can be "max" or "min"
+    verbose: True
+    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+    filename: '{epoch:02d}-{valid_avg_acc_epoch:.3f}'
+  learning_rate_monitor:
+    _target_: pytorch_lightning.callbacks.LearningRateMonitor
+    logging_interval: "step"
+  #audio_logging:
+  #  _target_: remfx.callbacks.AudioCallback
+  #  sample_rate: ${sample_rate}
+  #  log_audio: ${log_audio}
+
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  precision: 32 # Precision used for tensors, default `32`
+  min_epochs: 0
+  max_epochs: 300
+  log_every_n_steps: 1 # Logs metrics every N batches
+  accumulate_grad_batches: 1
+  accelerator: ${accelerator}
+  devices: 1
+  gradient_clip_val: 10.0
+  max_steps: -1

cfg/exp/5-5_cls_dynamic.yaml

@@ -0,0 +1,111 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "/scratch/cjs-logs"
+render_files: True
+render_root: "/scratch/EffectSet_cjs"
+accelerator: "gpu"
+log_audio: False
+# 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:
+  _target_: remfx.datasets.EffectDatamodule
+  train_dataset:
+    _target_: remfx.datasets.DynamicEffectDataset
+    total_chunks: 8000
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size: ${chunk_size}
+    mode: "train"
+    effect_modules: ${effects}
+    effects_to_keep: ${effects_to_keep}
+    effects_to_remove: ${effects_to_remove}
+    num_kept_effects: ${num_kept_effects}
+    num_removed_effects: ${num_removed_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
+    render_files: ${render_files}
+    render_root: ${render_root}
+    parallel: True
+  val_dataset:
+    _target_: remfx.datasets.EffectDataset
+    total_chunks: 1000
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size: ${chunk_size}
+    mode: "val"
+    effect_modules: ${effects}
+    effects_to_keep: ${effects_to_keep}
+    effects_to_remove: ${effects_to_remove}
+    num_kept_effects: ${num_kept_effects}
+    num_removed_effects: ${num_removed_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
+    render_files: ${render_files}
+    render_root: ${render_root}
+  test_dataset:
+    _target_: remfx.datasets.EffectDataset
+    total_chunks: 1000
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    chunk_size: ${chunk_size}
+    mode: "test"
+    effect_modules: ${effects}
+    effects_to_keep: ${effects_to_keep}
+    effects_to_remove: ${effects_to_remove}
+    num_kept_effects: ${num_kept_effects}
+    num_removed_effects: ${num_removed_effects}
+    shuffle_kept_effects: ${shuffle_kept_effects}
+    shuffle_removed_effects: ${shuffle_removed_effects}
+    render_files: ${render_files}
+    render_root: ${render_root}
+  train_batch_size: 32
+  test_batch_size: 256
+  num_workers: 12
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: "valid_avg_acc_epoch"   # name of the logged metric which determines when model is improving
+    save_top_k: 1           # save k best models (determined by above metric)
+    save_last: True         # additionaly always save model from last epoch
+    mode: "max"             # can be "max" or "min"
+    verbose: True
+    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+    filename: '{epoch:02d}-{valid_avg_acc_epoch:.3f}'
+  learning_rate_monitor:
+    _target_: pytorch_lightning.callbacks.LearningRateMonitor
+    logging_interval: "step"
+  #audio_logging:
+  #  _target_: remfx.callbacks.AudioCallback
+  #  sample_rate: ${sample_rate}
+  #  log_audio: ${log_audio}
+
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  precision: 32 # Precision used for tensors, default `32`
+  min_epochs: 0
+  max_epochs: 300
+  log_every_n_steps: 1 # Logs metrics every N batches
+  accumulate_grad_batches: 1
+  accelerator: ${accelerator}
+  devices: 1
+  gradient_clip_val: 10.0
+  max_steps: -1

cfg/model/{classifier.yaml → cls_panns_16k.yaml}

@@ -1,14 +1,15 @@
 # @package _global_
 model:
   _target_: remfx.models.FXClassifier
-  lr: 1e-4
+  lr: 3e-4
   lr_weight_decay: 1e-3
   sample_rate: ${sample_rate}
   network:
-    _target_: remfx.cnn14.Cnn14
+    _target_: remfx.classifier.Cnn14
     num_classes: ${num_classes}
-    n_fft: 4096
+    n_fft: 2048
     hop_length: 512
     n_mels: 128
     sample_rate: ${sample_rate}
+    model_sample_rate: 16000
 

cfg/model/cls_panns_44k_label_smoothing.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  mixup: True
+  label_smoothing: 0.1
+  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

cfg/model/cls_panns_48k.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+model:
+  _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
+  

cfg/model/cls_panns_48k_64.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+model:
+  _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: 64
+    sample_rate: ${sample_rate}
+    model_sample_rate: ${sample_rate}
+    specaugment: False
+  

cfg/model/cls_panns_48k_mixup.yaml

@@ -0,0 +1,16 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  mixup: True
+  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

cfg/model/cls_panns_48k_specaugment.yaml

@@ -0,0 +1,16 @@
+# @package _global_
+model:
+  _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: True

cfg/model/cls_panns_48k_specaugment_label_smoothing.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  mixup: False
+  label_smoothing: 0.15
+  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: True

cfg/model/cls_panns_pt.yaml

@@ -0,0 +1,12 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  mixup: False
+  network:
+    _target_: remfx.classifier.PANNs
+    num_classes: ${num_classes}
+    sample_rate: ${sample_rate}
+

cfg/model/cls_vggish.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.VGGish
+    num_classes: ${num_classes}
+    sample_rate: ${sample_rate}
+

cfg/model/cls_wav2clip.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.Wav2CLIP
+    num_classes: ${num_classes}
+    sample_rate: ${sample_rate}
+

cfg/model/cls_wav2vec2.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.wav2vec2
+    num_classes: ${num_classes}
+    sample_rate: ${sample_rate}
+

remfx/{cnn14.py → classifier.py}

@@ -1,8 +1,132 @@
 import torch
 import torchaudio
 import torch.nn as nn
+import hearbaseline
+import hearbaseline.vggish
+import hearbaseline.wav2vec2
+
+import wav2clip_hear
+import panns_hear
+
+
 import torch.nn.functional as F
-from utils import init_bn, init_layer
+from remfx.utils import init_bn, init_layer
+
+
+class PANNs(torch.nn.Module):
+    def __init__(
+        self, num_classes: int, sample_rate: float, hidden_dim: int = 256
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.model = panns_hear.load_model("hear2021-panns_hear.pth")
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=sample_rate, new_freq=32000
+        )
+        self.proj = torch.nn.Sequential(
+            torch.nn.Linear(2048, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        with torch.no_grad():
+            x = self.resample(x)
+            embed = panns_hear.get_scene_embeddings(x.view(x.shape[0], -1), self.model)
+        return self.proj(embed)
+
+
+class Wav2CLIP(nn.Module):
+    def __init__(
+        self,
+        num_classes: int,
+        sample_rate: float,
+        hidden_dim: int = 256,
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.model = wav2clip_hear.load_model("")
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=sample_rate, new_freq=16000
+        )
+        self.proj = torch.nn.Sequential(
+            torch.nn.Linear(512, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        with torch.no_grad():
+            x = self.resample(x)
+            embed = wav2clip_hear.get_scene_embeddings(
+                x.view(x.shape[0], -1), self.model
+            )
+        return self.proj(embed)
+
+
+class VGGish(nn.Module):
+    def __init__(
+        self,
+        num_classes: int,
+        sample_rate: float,
+        hidden_dim: int = 256,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=sample_rate, new_freq=16000
+        )
+        self.model = hearbaseline.vggish.load_model()
+        self.proj = torch.nn.Sequential(
+            torch.nn.Linear(128, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        with torch.no_grad():
+            x = self.resample(x)
+            embed = hearbaseline.vggish.get_scene_embeddings(
+                x.view(x.shape[0], -1), self.model
+            )
+        return self.proj(embed)
+
+
+class wav2vec2(nn.Module):
+    def __init__(
+        self,
+        num_classes: int,
+        sample_rate: float,
+        hidden_dim: int = 256,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=sample_rate, new_freq=16000
+        )
+        self.model = hearbaseline.wav2vec2.load_model()
+        self.proj = torch.nn.Sequential(
+            torch.nn.Linear(1024, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        with torch.no_grad():
+            x = self.resample(x)
+            embed = hearbaseline.wav2vec2.get_scene_embeddings(
+                x.view(x.shape[0], -1), self.model
+            )
+        return self.proj(embed)
+
 
 # adapted from https://github.com/qiuqiangkong/audioset_tagging_cnn/blob/master/pytorch/models.py
 
@@ -12,20 +136,25 @@ class Cnn14(nn.Module):
         self,
         num_classes: int,
         sample_rate: float,
-        n_fft: int = 2048,
-        hop_length: int = 512,
+        model_sample_rate: float,
+        n_fft: int = 1024,
+        hop_length: int = 256,
         n_mels: int = 128,
+        specaugment: bool = False,
     ):
         super().__init__()
         self.num_classes = num_classes
         self.n_fft = n_fft
         self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.model_sample_rate = model_sample_rate
+        self.specaugment = specaugment
 
         window = torch.hann_window(n_fft)
         self.register_buffer("window", window)
 
         self.melspec = torchaudio.transforms.MelSpectrogram(
-            sample_rate,
+            model_sample_rate,
             n_fft,
             hop_length=hop_length,
             n_mels=n_mels,
@@ -41,50 +170,80 @@ class Cnn14(nn.Module):
         self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)
 
         self.fc1 = nn.Linear(2048, 2048, bias=True)
-        self.fc_audioset = nn.Linear(2048, num_classes, bias=True)
+
+        # self.fc_audioset = nn.Linear(2048, num_classes, bias=True)
+        self.heads = torch.nn.ModuleList()
+        for _ in range(num_classes):
+            self.heads.append(nn.Linear(2048, 1, bias=True))
 
         self.init_weight()
 
+        if sample_rate != model_sample_rate:
+            self.resample = torchaudio.transforms.Resample(
+                orig_freq=sample_rate, new_freq=model_sample_rate
+            )
+
+        if self.specaugment:
+            self.freq_mask = torchaudio.transforms.FrequencyMasking(64, True)
+            self.time_mask = torchaudio.transforms.TimeMasking(128, True)
+
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
-        init_layer(self.fc_audioset)
+        # init_layer(self.fc_audioset)
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor, train: bool = False):
         """
         Input: (batch_size, data_length)"""
 
+        if self.sample_rate != self.model_sample_rate:
+            x = self.resample(x)
+
         x = self.melspec(x)
-        x = x.permute(0, 2, 1, 3)
-        x = self.bn0(x)
-        x = x.permute(0, 2, 1, 3)
 
-        if self.training:
-            pass
-            # x = self.spec_augmenter(x)
+        if self.specaugment and train:
+            # import matplotlib.pyplot as plt
+            # fig, axs = plt.subplots(2, 1, sharex=True)
+            # axs[0].imshow(x[0, :, :, :].detach().squeeze().cpu().numpy())
+            x = self.freq_mask(x)
+            x = self.time_mask(x)
+            # axs[1].imshow(x[0, :, :, :].detach().squeeze().cpu().numpy())
+            # plt.savefig("spec_augment.png", dpi=300)
+
+        # x = x.permute(0, 2, 1, 3)
+        # x = self.bn0(x)
+        # x = x.permute(0, 2, 1, 3)
+
+        # apply standardization
+        x = (x - x.mean(dim=0, keepdim=True)) / x.std(dim=0, keepdim=True)
 
         x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg")
-        x = F.dropout(x, p=0.2, training=self.training)
+        x = F.dropout(x, p=0.2, training=train)
         x = torch.mean(x, dim=3)
 
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
-        x = F.dropout(x, p=0.5, training=self.training)
+        x = F.dropout(x, p=0.5, training=train)
         x = F.relu_(self.fc1(x))
-        clipwise_output = self.fc_audioset(x)
 
-        return clipwise_output
+        outputs = []
+        for head in self.heads:
+            outputs.append(torch.sigmoid(head(x)))
+
+        # clipwise_output = self.fc_audioset(x)
+
+        return outputs
 
 
 class ConvBlock(nn.Module):

remfx/datasets.py

@@ -8,15 +8,16 @@ import pytorch_lightning as pl
 import random
 from tqdm import tqdm
 from pathlib import Path
-from remfx import effects
+from remfx import effects as effect_lib
 from typing import Any, List, Dict
 from torch.utils.data import Dataset, DataLoader
 from remfx.utils import select_random_chunk
+import multiprocessing
 
 
 # https://zenodo.org/record/1193957 -> VocalSet
 
-ALL_EFFECTS = effects.Pedalboard_Effects
+ALL_EFFECTS = effect_lib.Pedalboard_Effects
 # print(ALL_EFFECTS)
 
 
@@ -146,6 +147,230 @@ def locate_files(root: str, mode: str):
     return file_list
 
 
+def parallel_process_effects(
+    chunk_idx: int,
+    proc_root: str,
+    files: list,
+    chunk_size: int,
+    effects: list,
+    effects_to_keep: list,
+    num_kept_effects: tuple,
+    shuffle_kept_effects: bool,
+    effects_to_remove: list,
+    num_removed_effects: tuple,
+    shuffle_removed_effects: bool,
+    sample_rate: int,
+    target_lufs_db: float,
+):
+    """Note: This function has an issue with random seed. It may not fully randomize the effects."""
+    chunk = None
+    random_dataset_choice = random.choice(files)
+    while chunk is None:
+        random_file_choice = random.choice(random_dataset_choice)
+        chunk = select_random_chunk(random_file_choice, chunk_size, sample_rate)
+
+    # Sum to mono
+    if chunk.shape[0] > 1:
+        chunk = chunk.sum(0, keepdim=True)
+
+    dry = chunk
+
+    # loudness normalization
+    normalize = effect_lib.LoudnessNormalize(sample_rate, target_lufs_db=target_lufs_db)
+
+    # Apply Kept Effects
+    # Shuffle effects if specified
+    if shuffle_kept_effects:
+        effect_indices = torch.randperm(len(effects_to_keep))
+    else:
+        effect_indices = torch.arange(len(effects_to_keep))
+
+    r1 = num_kept_effects[0]
+    r2 = num_kept_effects[1]
+    num_kept_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
+    effect_indices = effect_indices[:num_kept_effects]
+    # Index in effect settings
+    effect_names_to_apply = [effects_to_keep[i] for i in effect_indices]
+    effects_to_apply = [effects[i] for i in effect_names_to_apply]
+    # Apply
+    dry_labels = []
+    for effect in effects_to_apply:
+        # Normalize in-between effects
+        dry = normalize(effect(dry))
+        dry_labels.append(ALL_EFFECTS.index(type(effect)))
+
+    # Apply effects_to_remove
+    # Shuffle effects if specified
+    if shuffle_removed_effects:
+        effect_indices = torch.randperm(len(effects_to_remove))
+    else:
+        effect_indices = torch.arange(len(effects_to_remove))
+    wet = torch.clone(dry)
+    r1 = num_removed_effects[0]
+    r2 = 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 = [effects_to_remove[i] for i in effect_indices]
+    effects_to_apply = [effects[i] for i in effect_names_to_apply]
+    # Apply
+    wet_labels = []
+    for effect in effects_to_apply:
+        # Normalize in-between effects
+        wet = 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 = normalize(dry)
+    normalized_wet = normalize(wet)
+
+    output_dir = proc_root / str(chunk_idx)
+    output_dir.mkdir(exist_ok=True)
+    torchaudio.save(output_dir / "input.wav", normalized_wet, sample_rate)
+    torchaudio.save(output_dir / "target.wav", normalized_dry, sample_rate)
+    torch.save(dry_labels_tensor, output_dir / "dry_effects.pt")
+    torch.save(wet_labels_tensor, output_dir / "wet_effects.pt")
+
+    # return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
+
+
+class DynamicEffectDataset(Dataset):
+    def __init__(
+        self,
+        root: str,
+        sample_rate: int,
+        chunk_size: int = 262144,
+        total_chunks: int = 1000,
+        effect_modules: List[Dict[str, torch.nn.Module]] = None,
+        effects_to_keep: List[str] = None,
+        effects_to_remove: List[str] = None,
+        num_kept_effects: List[int] = [1, 5],
+        num_removed_effects: List[int] = [1, 5],
+        shuffle_kept_effects: bool = True,
+        shuffle_removed_effects: bool = False,
+        render_files: bool = True,
+        render_root: str = None,
+        mode: str = "train",
+        parallel: bool = False,
+    ) -> None:
+        super().__init__()
+        self.chunks = []
+        self.song_idx = []
+        self.root = Path(root)
+        self.render_root = Path(render_root)
+        self.chunk_size = chunk_size
+        self.total_chunks = total_chunks
+        self.sample_rate = sample_rate
+        self.mode = mode
+        self.num_kept_effects = num_kept_effects
+        self.num_removed_effects = num_removed_effects
+        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.effects = effect_modules
+        self.shuffle_kept_effects = shuffle_kept_effects
+        self.shuffle_removed_effects = shuffle_removed_effects
+        effects_string = "_".join(
+            self.effects_to_keep
+            + ["_"]
+            + self.effects_to_remove
+            + ["_"]
+            + [str(x) for x in num_kept_effects]
+            + ["_"]
+            + [str(x) for x in num_removed_effects]
+        )
+        # self.validate_effect_input()
+        # self.proc_root = self.render_root / "processed" / effects_string / self.mode
+        self.parallel = parallel
+        self.files = locate_files(self.root, self.mode)
+
+    def process_effects(self, dry: torch.Tensor):
+        # Apply Kept Effects
+        # Shuffle effects if specified
+        if self.shuffle_kept_effects:
+            effect_indices = torch.randperm(len(self.effects_to_keep))
+        else:
+            effect_indices = torch.arange(len(self.effects_to_keep))
+
+        r1 = self.num_kept_effects[0]
+        r2 = self.num_kept_effects[1]
+        num_kept_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
+        effect_indices = effect_indices[:num_kept_effects]
+        # 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)
+        return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
+
+    def __len__(self):
+        return self.total_chunks
+
+    def __getitem__(self, _: int):
+        chunk = None
+        random_dataset_choice = random.choice(self.files)
+        while chunk is None:
+            random_file_choice = random.choice(random_dataset_choice)
+            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)
+
+        dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
+
+        return wet, dry, dry_effects, wet_effects
+
+
 class EffectDataset(Dataset):
     def __init__(
         self,
@@ -163,6 +388,7 @@ class EffectDataset(Dataset):
         render_files: bool = True,
         render_root: str = None,
         mode: str = "train",
+        parallel: bool = False,
     ):
         super().__init__()
         self.chunks = []
@@ -177,7 +403,7 @@ class EffectDataset(Dataset):
         self.num_removed_effects = num_removed_effects
         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 = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
+        self.normalize = effect_lib.LoudnessNormalize(sample_rate, target_lufs_db=-20)
         self.effects = effect_modules
         self.shuffle_kept_effects = shuffle_kept_effects
         self.shuffle_removed_effects = shuffle_removed_effects
@@ -192,6 +418,7 @@ class EffectDataset(Dataset):
         )
         self.validate_effect_input()
         self.proc_root = self.render_root / "processed" / effects_string / self.mode
+        self.parallel = parallel
 
         self.files = locate_files(self.root, self.mode)
 
@@ -212,26 +439,50 @@ class EffectDataset(Dataset):
         if render_files:
             # Split audio file into chunks, resample, then apply random effects
             self.proc_root.mkdir(parents=True, exist_ok=True)
-            for num_chunk in tqdm(range(self.total_chunks)):
-                chunk = None
-                random_dataset_choice = random.choice(self.files)
-                while chunk is None:
-                    random_file_choice = random.choice(random_dataset_choice)
-                    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)
 
-                dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
-                output_dir = self.proc_root / str(num_chunk)
-                output_dir.mkdir(exist_ok=True)
-                torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
-                torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
-                torch.save(dry_effects, output_dir / "dry_effects.pt")
-                torch.save(wet_effects, output_dir / "wet_effects.pt")
+            if self.parallel:
+                items = [
+                    (
+                        chunk_idx,
+                        self.proc_root,
+                        self.files,
+                        self.chunk_size,
+                        self.effects,
+                        self.effects_to_keep,
+                        self.num_kept_effects,
+                        self.shuffle_kept_effects,
+                        self.effects_to_remove,
+                        self.num_removed_effects,
+                        self.shuffle_removed_effects,
+                        self.sample_rate,
+                        -20.0,
+                    )
+                    for chunk_idx in range(self.total_chunks)
+                ]
+                with multiprocessing.Pool(processes=32) as pool:
+                    pool.starmap(parallel_process_effects, items)
+                print(f"Done proccessing {self.total_chunks}", flush=True)
+            else:
+                for num_chunk in tqdm(range(self.total_chunks)):
+                    chunk = None
+                    random_dataset_choice = random.choice(self.files)
+                    while chunk is None:
+                        random_file_choice = random.choice(random_dataset_choice)
+                        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)
+
+                    dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
+                    output_dir = self.proc_root / str(num_chunk)
+                    output_dir.mkdir(exist_ok=True)
+                    torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
+                    torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
+                    torch.save(dry_effects, output_dir / "dry_effects.pt")
+                    torch.save(wet_effects, output_dir / "wet_effects.pt")
 
             print("Finished rendering")
         else:
@@ -402,7 +653,8 @@ class EffectDatamodule(pl.LightningDataModule):
         val_dataset,
         test_dataset,
         *,
-        batch_size: int,
+        train_batch_size: int,
+        test_batch_size: int,
         num_workers: int,
         pin_memory: bool = False,
         **kwargs: int,
@@ -411,7 +663,8 @@ class EffectDatamodule(pl.LightningDataModule):
         self.train_dataset = train_dataset
         self.val_dataset = val_dataset
         self.test_dataset = test_dataset
-        self.batch_size = batch_size
+        self.train_batch_size = train_batch_size
+        self.test_batch_size = test_batch_size
         self.num_workers = num_workers
         self.pin_memory = pin_memory
 
@@ -421,7 +674,7 @@ class EffectDatamodule(pl.LightningDataModule):
     def train_dataloader(self) -> DataLoader:
         return DataLoader(
             dataset=self.train_dataset,
-            batch_size=self.batch_size,
+            batch_size=self.train_batch_size,
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             shuffle=True,
@@ -430,7 +683,7 @@ class EffectDatamodule(pl.LightningDataModule):
     def val_dataloader(self) -> DataLoader:
         return DataLoader(
             dataset=self.val_dataset,
-            batch_size=self.batch_size,
+            batch_size=self.train_batch_size,
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             shuffle=False,
@@ -439,7 +692,7 @@ class EffectDatamodule(pl.LightningDataModule):
     def test_dataloader(self) -> DataLoader:
         return DataLoader(
             dataset=self.test_dataset,
-            batch_size=2,  # Use small, consistent batch size for testing
+            batch_size=self.test_batch_size,  # Use small, consistent batch size for testing
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             shuffle=False,

remfx/models.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 import torchmetrics
 import pytorch_lightning as pl
 from torch import Tensor, nn
@@ -409,6 +410,37 @@ class TCNModel(nn.Module):
         return output
 
 
+def mixup(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0):
+    """Mixup data augmentation for time-domain signals.
+    Args:
+        x (torch.Tensor): Batch of time-domain signals, shape [batch, 1, time].
+        y (torch.Tensor): Batch of labels, shape [batch, n_classes].
+        alpha (float): Beta distribution parameter.
+    Returns:
+        torch.Tensor: Mixed time-domain signals, shape [batch, 1, time].
+        torch.Tensor: Mixed labels, shape [batch, n_classes].
+        torch.Tensor: Lambda
+    """
+    batch_size = x.size(0)
+    if alpha > 0:
+        # lam = np.random.beta(alpha, alpha)
+        lam = np.random.uniform(0.25, 0.75, batch_size)
+        lam = torch.from_numpy(lam).float().to(x.device).view(batch_size, 1, 1)
+    else:
+        lam = 1
+
+    print(lam)
+    if np.random.rand() > 0.5:
+        index = torch.randperm(batch_size).to(x.device)
+        mixed_x = lam * x + (1 - lam) * x[index, :]
+        mixed_y = torch.logical_or(y, y[index, :]).float()
+    else:
+        mixed_x = x
+        mixed_y = y
+
+    return mixed_x, mixed_y, lam
+
+
 class FXClassifier(pl.LightningModule):
     def __init__(
         self,
@@ -416,20 +448,85 @@ class FXClassifier(pl.LightningModule):
         lr_weight_decay: float,
         sample_rate: float,
         network: nn.Module,
+        mixup: bool = False,
+        label_smoothing: float = 0.0,
     ):
         super().__init__()
         self.lr = lr
         self.lr_weight_decay = lr_weight_decay
         self.sample_rate = sample_rate
         self.network = network
+        self.effects = ["Reverb", "Chorus", "Delay", "Distortion", "Compressor"]
+        self.mixup = mixup
+        self.label_smoothing = label_smoothing
+
+        self.loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+        self.loss_fn = torch.nn.BCELoss()
+
+        if False:
+            self.train_f1 = torchmetrics.classification.MultilabelF1Score(
+                5, average="none", multidim_average="global"
+            )
+            self.val_f1 = torchmetrics.classification.MultilabelF1Score(
+                5, average="none", multidim_average="global"
+            )
+            self.test_f1 = torchmetrics.classification.MultilabelF1Score(
+                5, average="none", multidim_average="global"
+            )
+
+            self.train_f1_avg = torchmetrics.classification.MultilabelF1Score(
+                5, threshold=0.5, average="macro", multidim_average="global"
+            )
+            self.val_f1_avg = torchmetrics.classification.MultilabelF1Score(
+                5, threshold=0.5, average="macro", multidim_average="global"
+            )
+            self.test_f1_avg = torchmetrics.classification.MultilabelF1Score(
+                5, threshold=0.5, average="macro", multidim_average="global"
+            )
+
+            self.metrics = {
+                "train": self.train_acc,
+                "valid": self.val_acc,
+                "test": self.test_acc,
+            }
+
+            self.avg_metrics = {
+                "train": self.train_f1_avg,
+                "valid": self.val_f1_avg,
+                "test": self.test_f1_avg,
+            }
 
-    def forward(self, x: torch.Tensor):
-        return self.network(x)
+        self.metrics = torch.nn.ModuleDict()
+        for effect in self.effects:
+            self.metrics[f"train_{effect}_acc"] = torchmetrics.classification.Accuracy(
+                task="binary"
+            )
+            self.metrics[f"valid_{effect}_acc"] = torchmetrics.classification.Accuracy(
+                task="binary"
+            )
+            self.metrics[f"test_{effect}_acc"] = torchmetrics.classification.Accuracy(
+                task="binary"
+            )
+
+    def forward(self, x: torch.Tensor, train: bool = False):
+        return self.network(x, train=train)
 
     def common_step(self, batch, batch_idx, mode: str = "train"):
+        train = True if mode == "train" else False
         x, y, dry_label, wet_label = batch
-        pred_label = self.network(x)
-        loss = nn.functional.cross_entropy(pred_label, dry_label)
+
+        if mode == "train" and self.mixup:
+            x_mixed, label_mixed, lam = mixup(x, wet_label)
+            outputs = self(x_mixed, train)
+            loss = 0
+            for idx, output in enumerate(outputs):
+                loss += self.loss_fn(output.squeeze(-1), label_mixed[..., idx])
+        else:
+            outputs = self(x, train)
+            loss = 0
+            for idx, output in enumerate(outputs):
+                loss += self.loss_fn(output.squeeze(-1), wet_label[..., idx])
+
         self.log(
             f"{mode}_loss",
             loss,
@@ -440,11 +537,25 @@ class FXClassifier(pl.LightningModule):
             sync_dist=True,
         )
 
+        acc_metrics = []
+        for idx, effect_name in enumerate(self.effects):
+            acc_metric = self.metrics[f"{mode}_{effect_name}_acc"](
+                outputs[idx].squeeze(-1), wet_label[..., idx]
+            )
+            self.log(
+                f"{mode}_{effect_name}_acc",
+                acc_metric,
+                on_step=True,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            acc_metrics.append(acc_metric)
+
         self.log(
-            f"{mode}_mAP",
-            torchmetrics.functional.retrieval_average_precision(
-                pred_label, dry_label.long()
-            ),
+            f"{mode}_avg_acc",
+            torch.mean(torch.stack(acc_metrics)),
             on_step=True,
             on_epoch=True,
             prog_bar=True,

setup.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 from setuptools import setup, find_packages
 
-NAME = "REMFX"
-DESCRIPTION = ""
+NAME = "remfx"
+DESCRIPTION = "Universal audio effect removal"
 URL = ""
 EMAIL = "[email protected]"
 AUTHOR = "Matthew Rice"

train_all.sh

@@ -0,0 +1,6 @@
+python scripts/train.py +exp=5-5_cls.yaml model=cls_wav2vec2 render_files=False logs_dir=/scratch/cjs-log
+python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_44k render_files=False logs_dir=/scratch/cjs-log
+python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_16k render_files=False logs_dir=/scratch/cjs-log 
+python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_pt render_files=False logs_dir=/scratch/cjs-log
+python scripts/train.py +exp=5-5_cls.yaml model=cls_vggish render_files=False logs_dir=/scratch/cjs-log 
+python scripts/train.py +exp=5-5_cls.yaml model=cls_wav2clip render_files=False logs_dir=/scratch/cjs-log