Merge branch 'cjs--classifier-v2'

README.md

@@ -81,4 +81,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/exp/5-5_cls.yaml

@@ -0,0 +1,59 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./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:
+  batch_size: 64
+  num_workers: 8
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: "valid_f1_avg_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_f1_avg_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: -1
+  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: 80000

cfg/model/cls_panns_16k.yaml

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

cfg/model/cls_panns_44k.yaml

@@ -0,0 +1,15 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.Cnn14
+    num_classes: ${num_classes}
+    n_fft: 1024
+    hop_length: 256
+    n_mels: 128
+    sample_rate: 44100
+    model_sample_rate: 44100
+    specaugment: True

cfg/model/cls_panns_44k_noaug.yaml

@@ -0,0 +1,15 @@
+# @package _global_
+model:
+  _target_: remfx.models.FXClassifier
+  lr: 3e-4
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.classifier.Cnn14
+    num_classes: ${num_classes}
+    n_fft: 1024
+    hop_length: 256
+    n_mels: 128
+    sample_rate: 44100
+    model_sample_rate: 44100
+    specaugment: False

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

@@ -1,14 +1,11 @@
 # @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.PANNs
     num_classes: ${num_classes}
-    n_fft: 4096
-    hop_length: 512
-    n_mels: 128
     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):
+        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):
+        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):
+        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):
+        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,
@@ -45,42 +174,56 @@ class Cnn14(nn.Module):
 
         self.init_weight()
 
+        if sample_rate != model_sample_rate:
+            self.resample = torchaudio.transforms.Resample(
+                orig_freq=sample_rate, new_freq=model_sample_rate
+            )
+
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         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)
+
+        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)
 
-        if self.training:
-            pass
-            # x = self.spec_augmenter(x)
-
         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)
 

remfx/models.py

@@ -422,14 +422,32 @@ class FXClassifier(pl.LightningModule):
         self.lr_weight_decay = lr_weight_decay
         self.sample_rate = sample_rate
         self.network = network
+        self.effects = ["distortion", "compressor", "reverb", "chorus", "delay"]
 
-    def forward(self, x: torch.Tensor):
+        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.metrics = {
+            "train": self.train_f1,
+            "valid": self.val_f1,
+            "test": self.test_f1,
+        }
+
+    def forward(self, x: torch.Tensor, train: bool = False):
         return self.network(x)
 
     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)
+        pred_label = self(x, train)
+        loss = nn.functional.cross_entropy(pred_label, wet_label)
         self.log(
             f"{mode}_loss",
             loss,
@@ -440,11 +458,12 @@ 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}_mAP",
-            torchmetrics.functional.retrieval_average_precision(
-                pred_label, dry_label.long()
-            ),
+            f"{mode}_f1_avg",
+            avg_metrics,
             on_step=True,
             on_epoch=True,
             prog_bar=True,
@@ -452,6 +471,17 @@ class FXClassifier(pl.LightningModule):
             sync_dist=True,
         )
 
+        for idx, effect_name in enumerate(self.effects):
+            self.log(
+                f"{mode}_f1_{effect_name}",
+                metrics[idx],
+                on_step=True,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+
         return loss
 
     def training_step(self, batch, batch_idx):

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