changing linear layers to MLP

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/exp/5-5_cls.yaml

@@ -56,4 +56,4 @@ trainer:
   accelerator: ${accelerator}
   devices: 1
   gradient_clip_val: 10.0
-  max_steps: 150000
+  max_steps: 80000

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

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):