Added intermediate logging of chain_inference

cfg/config.yaml

@@ -117,7 +117,6 @@ trainer:
   precision: 32 # Precision used for tensors, default `32`
   min_epochs: 0
   max_epochs: -1
-  enable_model_summary: False
   log_every_n_steps: 1 # Logs metrics every N batches
   accumulate_grad_batches: 1
   accelerator: ${accelerator}

cfg/model/demucs.yaml

@@ -13,3 +13,5 @@ model:
     audio_channels: 1
     nfft: 4096
     sample_rate: ${sample_rate}
+    channels: 64
+

remfx/models.py

@@ -11,6 +11,7 @@ from umx.openunmix.model import OpenUnmix, Separator
 from remfx.utils import FADLoss, spectrogram
 from remfx.tcn import TCN
 from remfx.utils import causal_crop
+from remfx.callbacks import log_wandb_audio_batch
 from remfx import effects
 import asteroid
 
@@ -42,12 +43,26 @@ class RemFXChainInference(pl.LightningModule):
         ]
         output = []
         with torch.no_grad():
-            for elem, effect_chain in zip(x, effects):
+            for i, (elem, effect_chain) in enumerate(zip(x, effects)):
                 elem = elem.unsqueeze(0)  # Add batch dim
+                log_wandb_audio_batch(
+                    logger=self.logger,
+                    id=f"{i}_Before",
+                    samples=elem.cpu(),
+                    sampling_rate=self.sample_rate,
+                    caption=effect_chain,
+                )
                 for effect in effect_chain:
                     # Get correct model based on effect name. This is a bit hacky
                     # Then sample the model
                     elem = self.model[effect.__name__].model.sample(elem)
+                    log_wandb_audio_batch(
+                        logger=self.logger,
+                        id=f"{i}_{effect}",
+                        samples=elem.cpu(),
+                        sampling_rate=self.sample_rate,
+                        caption=effect_chain,
+                    )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
 

scripts/chain_inference.py

@@ -2,7 +2,6 @@ import pytorch_lightning as pl
 import hydra
 from omegaconf import DictConfig
 import remfx.utils as utils
-from pytorch_lightning.utilities.model_summary import ModelSummary
 import torch
 from remfx.models import RemFXChainInference
 
@@ -21,10 +20,9 @@ def main(cfg: DictConfig):
     for effect in cfg.ckpts:
         ckpt_path = cfg.ckpts[effect]
         model = hydra.utils.instantiate(cfg.model, _convert_="partial")
-        state_dict = torch.load(ckpt_path, map_location=torch.device("cpu"))[
-            "state_dict"
-        ]
+        state_dict = torch.load(ckpt_path)["state_dict"]
         model.load_state_dict(state_dict)
+        model.to(cfg.device)
         models[effect] = model
 
     callbacks = []
@@ -48,8 +46,6 @@ def main(cfg: DictConfig):
         callbacks=callbacks,
         logger=logger,
     )
-    summary = ModelSummary(model)
-    print(summary)
 
     inference_model = RemFXChainInference(
         models, sample_rate=cfg.sample_rate, num_bins=cfg.num_bins

scripts/test.py

@@ -44,8 +44,6 @@ def main(cfg: DictConfig):
         callbacks=callbacks,
         logger=logger,
     )
-    summary = ModelSummary(model)
-    print(summary)
     trainer.test(model=model, datamodule=datamodule)
 
 

scripts/train.py

@@ -2,7 +2,6 @@ import pytorch_lightning as pl
 import hydra
 from omegaconf import DictConfig
 import remfx.utils as utils
-from pytorch_lightning.utilities.model_summary import ModelSummary
 
 log = utils.get_logger(__name__)
 
@@ -39,8 +38,6 @@ def main(cfg: DictConfig):
         callbacks=callbacks,
         logger=logger,
     )
-    summary = ModelSummary(model)
-    print(summary)
     trainer.fit(model=model, datamodule=datamodule)
     trainer.test(model=model, datamodule=datamodule, ckpt_path="best")