Merge pull request #20 from mhrice/fix-demucs

README.md

@@ -10,10 +10,13 @@
 `./scripts/download_egfx.sh`
 
 ## Train model
-1. Change Wandb variables in `shell_vars.sh`
+1. Change Wandb variables in `shell_vars.sh` and `source shell_vars.sh`
 2. `python scripts/train.py exp=audio_diffusion`
 or
 2. `python scripts/train.py exp=umx`
+or
+2. `python scripts/train.py exp=demucs`
+
 
 To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-line
 

exp/demucs.yaml

@@ -9,7 +9,7 @@ model:
   sample_rate: ${sample_rate}
   network:
     _target_: remfx.models.DemucsModel
-    sources: ["other"]
+    sources: ["mixture"]
     audio_channels: 1
     nfft: 4096
     sample_rate: ${sample_rate}

remfx/models.py

@@ -38,6 +38,8 @@ class RemFXModel(pl.LightningModule):
                 "L1": L1Loss(),
             }
         )
+        # Log first batch metrics input vs output only once
+        self.log_first = True
 
     @property
     def device(self):
@@ -67,9 +69,14 @@ class RemFXModel(pl.LightningModule):
         x, y, label = batch
         # Metric logging
         for metric in self.metrics:
+            # SISDR returns negative values, so negate them
+            if metric == "SISDR":
+                negate = -1
+            else:
+                negate = 1
             self.log(
                 f"{mode}_{metric}",
-                self.metrics[metric](output, y),
+                negate * self.metrics[metric](output, y),
                 on_step=False,
                 on_epoch=True,
                 logger=True,
@@ -79,6 +86,26 @@ class RemFXModel(pl.LightningModule):
 
         return loss
 
+    def on_train_batch_start(self, batch, batch_idx):
+        if self.log_first:
+            x, target, label = batch
+            for metric in self.metrics:
+                # SISDR returns negative values, so negate them
+                if metric == "SISDR":
+                    negate = -1
+                else:
+                    negate = 1
+                self.log(
+                    f"Input_{metric}",
+                    negate * self.metrics[metric](x, target),
+                    on_step=False,
+                    on_epoch=True,
+                    logger=True,
+                    prog_bar=True,
+                    sync_dist=True,
+                )
+            self.log_first = False
+
     def on_validation_epoch_start(self):
         self.log_next = True