Merge branch 'main' into classifier-inference

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

@@ -5,9 +5,9 @@ defaults:
 seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
-logs_dir: "./logs"
+logs_dir: "/scratch/cjs-logs"
 render_files: True
-render_root: "/scratch/EffectSet_cjs_nobass"
+render_root: "/scratch/EffectSet_cjs"
 accelerator: "gpu"
 log_audio: False
 # Effects
@@ -24,19 +24,20 @@ effects_to_remove:
   - chorus
   - delay
 datamodule:
-  batch_size: 64
+  train_batch_size: 64
+  test_batch_size: 256
   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
+    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_f1_avg_epoch:.3f}'
+    filename: '{epoch:02d}-{valid_avg_acc_epoch:.3f}'
   learning_rate_monitor:
     _target_: pytorch_lightning.callbacks.LearningRateMonitor
     logging_interval: "step"
@@ -50,10 +51,10 @@ trainer:
   _target_: pytorch_lightning.Trainer
   precision: 32 # Precision used for tensors, default `32`
   min_epochs: 0
-  max_epochs: -1
+  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: 100000
+  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

remfx/classifier.py

@@ -172,7 +172,11 @@ 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()
 
@@ -188,7 +192,7 @@ class Cnn14(nn.Module):
     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, train: bool = False):
         """
@@ -208,9 +212,12 @@ class Cnn14(nn.Module):
             # 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)
+        # 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=train)
@@ -231,9 +238,14 @@ class Cnn14(nn.Module):
         x = x1 + x2
         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

@@ -162,6 +162,7 @@ def parallel_process_effects(
     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:
@@ -242,6 +243,134 @@ def parallel_process_effects(
     # 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,
@@ -530,7 +659,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,
@@ -539,7 +669,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
 
@@ -549,7 +680,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,
@@ -558,7 +689,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,

remfx/models.py

@@ -471,13 +471,20 @@ def mixup(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0):
     """
     batch_size = x.size(0)
     if alpha > 0:
-        lam = np.random.beta(alpha, alpha)
+        # 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
 
-    index = torch.randperm(batch_size).to(x.device)
-    mixed_x = lam * x + (1 - lam) * x[index, :]
-    mixed_y = lam * y + (1 - lam) * y[index, :]
+    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
 
@@ -502,38 +509,52 @@ class FXClassifier(pl.LightningModule):
         self.label_smoothing = label_smoothing
 
         self.loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+        self.loss_fn = torch.nn.BCELoss()
 
-        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"
-        )
+        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.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.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,
-        }
+            self.avg_metrics = {
+                "train": self.train_f1_avg,
+                "valid": self.val_f1_avg,
+                "test": self.test_f1_avg,
+            }
+
+        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)
@@ -544,15 +565,15 @@ class FXClassifier(pl.LightningModule):
 
         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, ...])
+            outputs = self(x_mixed, train)
+            loss = 0
+            for idx, output in enumerate(outputs):
+                loss += self.loss_fn(output.squeeze(-1), label_mixed[..., idx])
         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, ...])
+            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",
@@ -564,26 +585,25 @@ class FXClassifier(pl.LightningModule):
             sync_dist=True,
         )
 
-        metrics = self.metrics[mode](torch.sigmoid(pred_label), wet_label.long())
-
+        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}_f1_{effect_name}",
-                metrics[idx],
+                f"{mode}_{effect_name}_acc",
+                acc_metric,
                 on_step=True,
                 on_epoch=True,
                 prog_bar=True,
                 logger=True,
                 sync_dist=True,
             )
-
-        avg_metrics = self.avg_metrics[mode](
-            torch.sigmoid(pred_label), wet_label.long()
-        )
+            acc_metrics.append(acc_metric)
 
         self.log(
-            f"{mode}_f1_avg",
-            avg_metrics,
+            f"{mode}_avg_acc",
+            torch.mean(torch.stack(acc_metrics)),
             on_step=True,
             on_epoch=True,
             prog_bar=True,