Merge branch 'cjs--classifier-v2' of https://github.com/mhrice/RemFx into classifier-inference

cfg/exp/5-5_cls.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet_cjs"
+render_root: "/scratch/EffectSet_cjs_nobass"
 accelerator: "gpu"
 log_audio: False
 # Effects
@@ -56,4 +56,4 @@ trainer:
   accelerator: ${accelerator}
   devices: 1
   gradient_clip_val: 10.0
-  max_steps: 80000
+  max_steps: 100000

cfg/model/cls_panns_16k.yaml

@@ -10,6 +10,6 @@ model:
     n_fft: 2048
     hop_length: 512
     n_mels: 128
-    sample_rate: 44100
+    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_44k.yaml → cls_panns_48k_mixup.yaml}

@@ -4,12 +4,13 @@ model:
   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: 1024
-    hop_length: 256
+    n_fft: 2048
+    hop_length: 512
     n_mels: 128
-    sample_rate: 44100
-    model_sample_rate: 44100
-    specaugment: True
+    sample_rate: ${sample_rate}
+    model_sample_rate: ${sample_rate}
+    specaugment: False

cfg/model/{cls_panns_44k_noaug.yaml → cls_panns_48k_specaugment.yaml}

@@ -4,12 +4,13 @@ model:
   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: 1024
-    hop_length: 256
+    n_fft: 2048
+    hop_length: 512
     n_mels: 128
-    sample_rate: 44100
-    model_sample_rate: 44100
-    specaugment: False
+    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

@@ -4,6 +4,7 @@ model:
   lr: 3e-4
   lr_weight_decay: 1e-3
   sample_rate: ${sample_rate}
+  mixup: False
   network:
     _target_: remfx.classifier.PANNs
     num_classes: ${num_classes}

remfx/classifier.py

@@ -33,7 +33,7 @@ class PANNs(torch.nn.Module):
             torch.nn.Linear(hidden_dim, num_classes),
         )
 
-    def forward(self, x: torch.Tensor):
+    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)
@@ -61,7 +61,7 @@ class Wav2CLIP(nn.Module):
             torch.nn.Linear(hidden_dim, num_classes),
         )
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor, **kwargs):
         with torch.no_grad():
             x = self.resample(x)
             embed = wav2clip_hear.get_scene_embeddings(
@@ -91,7 +91,7 @@ class VGGish(nn.Module):
             torch.nn.Linear(hidden_dim, num_classes),
         )
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor, **kwargs):
         with torch.no_grad():
             x = self.resample(x)
             embed = hearbaseline.vggish.get_scene_embeddings(
@@ -121,7 +121,7 @@ class wav2vec2(nn.Module):
             torch.nn.Linear(hidden_dim, num_classes),
         )
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor, **kwargs):
         with torch.no_grad():
             x = self.resample(x)
             embed = hearbaseline.wav2vec2.get_scene_embeddings(
@@ -181,6 +181,10 @@ class Cnn14(nn.Module):
                 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)

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,101 @@ 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,
+):
+    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 EffectDataset(Dataset):
     def __init__(
         self,
@@ -163,6 +259,7 @@ class EffectDataset(Dataset):
         render_files: bool = True,
         render_root: str = None,
         mode: str = "train",
+        parallel: bool = True,
     ):
         super().__init__()
         self.chunks = []
@@ -177,7 +274,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 +289,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 +310,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:

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
@@ -424,6 +425,30 @@ 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)
+    else:
+        lam = 1
+
+    index = torch.randperm(batch_size).to(x.device)
+    mixed_x = lam * x + (1 - lam) * x[index, :]
+    mixed_y = lam * y + (1 - lam) * y[index, :]
+
+    return mixed_x, mixed_y, lam
+
+
 class FXClassifier(pl.LightningModule):
     def __init__(
         self,
@@ -431,13 +456,19 @@ 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 = ["distortion", "compressor", "reverb", "chorus", "delay"]
+        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.train_f1 = torchmetrics.classification.MultilabelF1Score(
             5, average="none", multidim_average="global"
@@ -449,20 +480,47 @@ class FXClassifier(pl.LightningModule):
             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_f1,
             "valid": self.val_f1,
             "test": self.test_f1,
         }
 
+        self.avg_metrics = {
+            "train": self.train_f1_avg,
+            "valid": self.val_f1_avg,
+            "test": self.test_f1_avg,
+        }
+
     def forward(self, x: torch.Tensor, train: bool = False):
-        return self.network(x)
+        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(x, train)
-        loss = nn.functional.cross_entropy(pred_label, wet_label)
+
+        if mode == "train" and self.mixup:
+            x_mixed, label_mixed, lam = mixup(x, wet_label)
+            pred_label = self(x_mixed, train)
+            loss = self.loss_fn(pred_label, label_mixed)
+            print(torch.sigmoid(pred_label[0, ...]))
+            print(label_mixed[0, ...])
+        else:
+            pred_label = self(x, train)
+            loss = self.loss_fn(pred_label, wet_label)
+            print(torch.where(torch.sigmoid(pred_label[0, ...]) > 0.5, 1.0, 0.0).long())
+            print(wet_label.long()[0, ...])
+
         self.log(
             f"{mode}_loss",
             loss,
@@ -473,18 +531,7 @@ class FXClassifier(pl.LightningModule):
             sync_dist=True,
         )
 
-        metrics = self.metrics[mode](pred_label, wet_label.long())
-        avg_metrics = torch.mean(metrics)
-
-        self.log(
-            f"{mode}_f1_avg",
-            avg_metrics,
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-            logger=True,
-            sync_dist=True,
-        )
+        metrics = self.metrics[mode](torch.sigmoid(pred_label), wet_label.long())
 
         for idx, effect_name in enumerate(self.effects):
             self.log(
@@ -497,6 +544,20 @@ class FXClassifier(pl.LightningModule):
                 sync_dist=True,
             )
 
+        avg_metrics = self.avg_metrics[mode](
+            torch.sigmoid(pred_label), wet_label.long()
+        )
+
+        self.log(
+            f"{mode}_f1_avg",
+            avg_metrics,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+
         return loss
 
     def training_step(self, batch, batch_idx):