Refactor to use hydra

.gitignore

@@ -6,4 +6,6 @@ data/
 .DS_Store
 __pycache__/
 lightning_logs/
-RemFX/
+RemFX/
+outputs/
+logs/

config.yaml

@@ -0,0 +1,50 @@
+defaults:
+  - _self_
+  - exp: null
+seed: 12345
+train: True
+length: 262144
+sample_rate: 22050
+logs_dir: "./logs"
+log_every_n_steps: 1000
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: "valid_loss"   # 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: "min"             # can be "max" or "min"
+    verbose: False
+    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+    filename: '{epoch:02d}-{valid_loss:.3f}'
+
+datamodule:
+  _target_: datasets.Datamodule
+  dataset:
+    _target_: datasets.GuitarFXDataset
+    sample_rate: ${sample_rate}
+    root: ${oc.env:DATASET_ROOT}
+    length: ${length}
+  val_split: 0.2
+  batch_size: 16
+  num_workers: 8
+  pin_memory: True
+
+logger:
+  _target_: pytorch_lightning.loggers.WandbLogger
+  project: ${oc.env:WANDB_PROJECT}
+  entity: ${oc.env:WANDB_ENTITY}
+  # offline: False  # set True to store all logs only locally
+  job_type: "train"
+  group: ""
+  save_dir: "."
+
+trainer:
+  _target_: pytorch_lightning.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

datasets.py

@@ -1,10 +1,10 @@
-import torch
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, DataLoader, random_split
 import torchaudio
 import torchaudio.transforms as T
 import torch.nn.functional as F
 from pathlib import Path
-from typing import List
+import pytorch_lightning as pl
+from typing import Any, List
 
 # https://zenodo.org/record/7044411/
 
@@ -18,18 +18,19 @@ class GuitarFXDataset(Dataset):
         root: str,
         sample_rate: int,
         length: int = LENGTH,
-        effect_type: List[str] = None,
+        effect_types: List[str] = None,
     ):
         self.length = length
         self.wet_files = []
         self.dry_files = []
         self.labels = []
         self.root = Path(root)
-        if effect_type is None:
-            effect_type = [
+
+        if effect_types is None:
+            effect_types = [
                 d.name for d in self.root.iterdir() if d.is_dir() and d != "Clean"
             ]
-        for i, effect in enumerate(effect_type):
+        for i, effect in enumerate(effect_types):
             for pickup in Path(self.root / effect).iterdir():
                 self.wet_files += sorted(list(pickup.glob("*.wav")))
                 self.dry_files += sorted(
@@ -61,3 +62,50 @@ class GuitarFXDataset(Dataset):
         elif resampled_y.shape[-1] > self.length:
             resampled_y = resampled_y[:, : self.length]
         return (resampled_x, resampled_y, effect_label)
+
+
+class Datamodule(pl.LightningDataModule):
+    def __init__(
+        self,
+        dataset,
+        *,
+        val_split: float,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool = False,
+        **kwargs: int,
+    ) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.val_split = val_split
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+        self.data_train: Any = None
+        self.data_val: Any = None
+
+    def setup(self, stage: Any = None) -> None:
+        split = [1.0 - self.val_split, self.val_split]
+        train_size = int(split[0] * len(self.dataset))
+        val_size = int(split[1] * len(self.dataset))
+        self.data_train, self.data_val = random_split(
+            self.dataset, [train_size, val_size]
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+            dataset=self.data_train,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+            dataset=self.data_val,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+        )

exp/audio_diffusion.yaml

@@ -0,0 +1,15 @@
+# @package _global_
+model: 
+  _target_: models.RemFXModel
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  network:
+    _target_: models.DiffusionGenerationModel
+    n_channels: 1
+datamodule:
+  dataset:
+    effect_types: ["Clean"]
+  batch_size: 2

exp/demucs.yaml

@@ -0,0 +1 @@
+# @package _global_

exp/umx.yaml

@@ -0,0 +1,18 @@
+# @package _global_
+model: 
+  _target_: models.RemFXModel
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  network:
+    _target_: models.OpenUnmixModel
+    n_fft: 2048
+    hop_length: 512
+    n_channels: 1
+    alpha: 0.3
+    sample_rate: ${sample_rate}
+datamodule:
+  dataset:
+    effect_types: ["RAT"]

main.py

@@ -1,19 +0,0 @@
-from audio_diffusion_pytorch import AudioDiffusionModel
-import torch
-from tqdm import tqdm
-import wandb
-
-model = AudioDiffusionModel(in_channels=1)
-wandb.init(project="RemFX", entity="mattricesound")
-
-x = torch.randn(2, 1, 2**18)
-for i in tqdm(range(100)):
-    loss = model(x)
-    loss.backward()
-    if i % 10 == 0:
-        print(loss)
-        wandb.log({"loss": loss})
-
-
-noise = torch.randn(2, 1, 2**18)
-sampled = model.sample(noise=noise, num_steps=5)

models.py

@@ -1,9 +1,10 @@
 import torch
-from torch import Tensor
+from torch import Tensor, nn
 import pytorch_lightning as pl
 from einops import rearrange
 import wandb
 from audio_diffusion_pytorch import AudioDiffusionModel
+import auraloss
 
 import sys
 
@@ -14,50 +15,49 @@ from umx.openunmix.model import OpenUnmix, Separator
 SAMPLE_RATE = 22050  # From audio-diffusion-pytorch
 
 
-class OpenUnmixModel(pl.LightningModule):
+class RemFXModel(pl.LightningModule):
     def __init__(
         self,
-        n_fft: int = 2048,
-        hop_length: int = 512,
-        alpha: float = 0.3,
+        lr: float,
+        lr_beta1: float,
+        lr_beta2: float,
+        lr_eps: float,
+        lr_weight_decay: float,
+        network: nn.Module,
     ):
         super().__init__()
-        self.model = OpenUnmix(
-            nb_channels=1,
-            nb_bins=n_fft // 2 + 1,
-        )
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.alpha = alpha
-        window = torch.hann_window(n_fft)
-        self.register_buffer("window", window)
+        self.lr = lr
+        self.lr_beta1 = lr_beta1
+        self.lr_beta2 = lr_beta2
+        self.lr_eps = lr_eps
+        self.lr_weight_decay = lr_weight_decay
+        self.model = network
 
-    def forward(self, x: torch.Tensor):
-        return self.model(x)
+    @property
+    def device(self):
+        return next(self.model.parameters()).device
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            list(self.model.parameters()),
+            lr=self.lr,
+            betas=(self.lr_beta1, self.lr_beta2),
+            eps=self.lr_eps,
+            weight_decay=self.lr_weight_decay,
+        )
+        return optimizer
 
     def training_step(self, batch, batch_idx):
-        loss, _ = self.common_step(batch, batch_idx, mode="train")
+        loss = self.common_step(batch, batch_idx, mode="train")
         return loss
 
     def validation_step(self, batch, batch_idx):
-        loss, Y = self.common_step(batch, batch_idx, mode="val")
-        return loss, Y
+        loss = self.common_step(batch, batch_idx, mode="valid")
 
     def common_step(self, batch, batch_idx, mode: str = "train"):
-        x, target, label = batch
-        X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
-        Y = self(X)
-        Y_hat = spectrogram(
-            target, self.window, self.n_fft, self.hop_length, self.alpha
-        )
-        loss = torch.nn.functional.mse_loss(Y, Y_hat)
-        self.log(f"{mode}_loss", loss, on_step=True, on_epoch=True)
-        return loss, Y
-
-    def configure_optimizers(self):
-        return torch.optim.Adam(
-            self.parameters(), lr=1e-4, betas=(0.95, 0.999), eps=1e-6, weight_decay=1e-3
-        )
+        loss = self.model(batch)
+        self.log(f"{mode}_loss", loss)
+        return loss
 
     def on_validation_epoch_start(self):
         self.log_next = True
@@ -65,14 +65,7 @@ class OpenUnmixModel(pl.LightningModule):
     def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
         if self.log_next:
             x, target, label = batch
-            s = Separator(
-                target_models={"other": self.model},
-                nb_channels=1,
-                sample_rate=SAMPLE_RATE,
-                n_fft=self.n_fft,
-                n_hop=self.hop_length,
-            ).to(self.device)
-            outputs = s(x).squeeze(1)
+            y = self.model.sample(x)
             log_wandb_audio_batch(
                 logger=self.logger,
                 id="sample",
@@ -83,12 +76,12 @@ class OpenUnmixModel(pl.LightningModule):
             log_wandb_audio_batch(
                 logger=self.logger,
                 id="prediction",
-                samples=outputs.cpu(),
+                samples=y.cpu(),
                 sampling_rate=SAMPLE_RATE,
                 caption=f"Epoch {self.current_epoch}",
             )
             log_wandb_audio_batch(
-                logger=self.loggger,
+                logger=self.logger,
                 id="target",
                 samples=target.cpu(),
                 sampling_rate=SAMPLE_RATE,
@@ -97,55 +90,65 @@ class OpenUnmixModel(pl.LightningModule):
             self.log_next = False
 
 
-class DiffusionGenerationModel(pl.LightningModule):
-    def __init__(self, model: torch.nn.Module):
+class OpenUnmixModel(torch.nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        n_channels: int = 1,
+        alpha: float = 0.3,
+        sample_rate: int = 22050,
+    ):
         super().__init__()
-        self.model = model
+        self.n_channels = n_channels
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.alpha = alpha
+        window = torch.hann_window(n_fft)
+        self.register_buffer("window", window)
 
-    def forward(self, x: torch.Tensor):
-        return self.model(x)
+        self.num_bins = self.n_fft // 2 + 1
+        self.sample_rate = sample_rate
+        self.model = OpenUnmix(
+            nb_channels=self.n_channels,
+            nb_bins=self.num_bins,
+        )
+        self.separator = Separator(
+            target_models={"other": self.model},
+            nb_channels=self.n_channels,
+            sample_rate=self.sample_rate,
+            n_fft=self.n_fft,
+            n_hop=self.hop_length,
+        )
+        self.loss_fn = auraloss.freq.MultiResolutionSTFTLoss(
+            n_bins=self.num_bins, sample_rate=self.sample_rate
+        )
 
-    def sample(self, *args, **kwargs) -> Tensor:
-        return self.model.sample(*args, **kwargs)
+    def forward(self, batch):
+        x, target, label = batch
+        X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
+        Y = self.model(X)
+        sep_out = self.separator(x).squeeze(1)
+        loss = self.loss_fn(sep_out, target)
 
-    def training_step(self, batch, batch_idx):
-        loss = self.common_step(batch, batch_idx, mode="train")
         return loss
 
-    def validation_step(self, batch, batch_idx):
-        loss = self.common_step(batch, batch_idx, mode="val")
+    def sample(self, x: Tensor) -> Tensor:
+        return self.separator(x).squeeze(1)
 
-    def common_step(self, batch, batch_idx, mode: str = "train"):
-        x, target, label = batch
-        loss = self(x)
-        self.log(f"{mode}_loss", loss, on_step=True, on_epoch=True)
-        return loss
 
-    def configure_optimizers(self):
-        return torch.optim.Adam(
-            self.parameters(), lr=1e-4, betas=(0.95, 0.999), eps=1e-6, weight_decay=1e-3
-        )
-
-    def on_validation_epoch_start(self):
-        self.log_next = True
+class DiffusionGenerationModel(nn.Module):
+    def __init__(self, n_channels: int = 1):
+        super().__init__()
+        self.model = AudioDiffusionModel(in_channels=n_channels)
 
-    def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
+    def forward(self, batch):
         x, target, label = batch
-        if self.log_next:
-            self.log_sample(x)
-            self.log_next = False
+        return self.model(x)
 
-    @torch.no_grad()
-    def log_sample(self, batch, num_steps=10):
-        # Get start diffusion noise
-        noise = torch.randn(batch.shape, device=self.device)
-        sampled = self.sample(noise=noise, num_steps=num_steps)  # Suggested range: 2-50
-        log_wandb_audio_batch(
-            id="sample",
-            samples=sampled,
-            sampling_rate=SAMPLE_RATE,
-            caption=f"Sampled in {num_steps} steps",
-        )
+    def sample(self, x: Tensor, num_steps: int = 10) -> Tensor:
+        noise = torch.randn(x.shape)
+        return self.model.sample(noise, num_steps=num_steps)
 
 
 def log_wandb_audio_batch(

setup.py

@@ -42,6 +42,8 @@ setup(
         "ema_pytorch",
         "einops",
         "librosa",
+        "hydra-core",
+        "auraloss",
     ],
     include_package_data=True,
     license="Apache License 2.0",

shell_vars.sh

@@ -0,0 +1,3 @@
+export DATASET_ROOT="/Users/matthewrice/Developer/remfx/data/egfx"
+export WANDB_PROJECT="RemFX"
+export WANDB_ENTITY="mattricesound"

train.py

@@ -1,35 +1,50 @@
 from pytorch_lightning.loggers import WandbLogger
 import pytorch_lightning as pl
-import torch
 from torch.utils.data import DataLoader
 from datasets import GuitarFXDataset
 from models import DiffusionGenerationModel, OpenUnmixModel
+import hydra
+from omegaconf import DictConfig
+import utils
 
+log = utils.get_logger(__name__)
 
-SAMPLE_RATE = 22050
-TRAIN_SPLIT = 0.8
 
+@hydra.main(version_base=None, config_path=".", config_name="config.yaml")
+def main(cfg: DictConfig):
+    # Apply seed for reproducibility
+    print(cfg)
+    pl.seed_everything(cfg.seed)
 
-def main():
-    wandb_logger = WandbLogger(project="RemFX", save_dir="./")
-    trainer = pl.Trainer(logger=wandb_logger, max_epochs=100)
-    guitfx = GuitarFXDataset(
-        root="./data/egfx",
-        sample_rate=SAMPLE_RATE,
-        effect_type=["Phaser"],
-    )
-    train_size = int(TRAIN_SPLIT * len(guitfx))
-    val_size = len(guitfx) - train_size
-    train_dataset, val_dataset = torch.utils.data.random_split(
-        guitfx, [train_size, val_size]
-    )
-    train = DataLoader(train_dataset, batch_size=2)
-    val = DataLoader(val_dataset, batch_size=2)
+    log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>.")
+    datamodule = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
+
+    log.info(f"Instantiating model <{cfg.model._target_}>.")
+    model = hydra.utils.instantiate(cfg.model, _convert_="partial")
 
-    # model = DiffusionGenerationModel()
-    model = OpenUnmixModel()
+    # Init all callbacks
+    callbacks = []
+    if "callbacks" in cfg:
+        for _, cb_conf in cfg["callbacks"].items():
+            if "_target_" in cb_conf:
+                log.info(f"Instantiating callback <{cb_conf._target_}>.")
+                callbacks.append(hydra.utils.instantiate(cb_conf, _convert_="partial"))
 
-    trainer.fit(model=model, train_dataloaders=train, val_dataloaders=val)
+    logger = hydra.utils.instantiate(cfg.logger, _convert_="partial")
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>.")
+    trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=logger, _convert_="partial"
+    )
+    log.info("Logging hyperparameters!")
+    utils.log_hyperparameters(
+        config=cfg,
+        model=model,
+        datamodule=datamodule,
+        trainer=trainer,
+        callbacks=callbacks,
+        logger=logger,
+    )
+    trainer.fit(model=model, datamodule=datamodule)
 
 
 if __name__ == "__main__":

utils.py

@@ -0,0 +1,71 @@
+import logging
+from typing import List
+import pytorch_lightning as pl
+from omegaconf import DictConfig
+from pytorch_lightning.utilities import rank_zero_only
+
+
+def get_logger(name=__name__) -> logging.Logger:
+    """Initializes multi-GPU-friendly python command line logger."""
+
+    logger = logging.getLogger(name)
+
+    # this ensures all logging levels get marked with the rank zero decorator
+    # otherwise logs would get multiplied for each GPU process in multi-GPU setup
+    for level in (
+        "debug",
+        "info",
+        "warning",
+        "error",
+        "exception",
+        "fatal",
+        "critical",
+    ):
+        setattr(logger, level, rank_zero_only(getattr(logger, level)))
+
+    return logger
+
+
+log = get_logger(__name__)
+
+
+@rank_zero_only
+def log_hyperparameters(
+    config: DictConfig,
+    model: pl.LightningModule,
+    datamodule: pl.LightningDataModule,
+    trainer: pl.Trainer,
+    callbacks: List[pl.Callback],
+    logger: pl.loggers.LightningLoggerBase,
+) -> None:
+    """Controls which config parts are saved by Lightning loggers.
+    Additionaly saves:
+    - number of model parameters
+    """
+
+    if not trainer.logger:
+        return
+
+    hparams = {}
+
+    # choose which parts of hydra config will be saved to loggers
+    hparams["model"] = config["model"]
+
+    # save number of model parameters
+    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
+    hparams["model/params/trainable"] = sum(
+        p.numel() for p in model.parameters() if p.requires_grad
+    )
+    hparams["model/params/non_trainable"] = sum(
+        p.numel() for p in model.parameters() if not p.requires_grad
+    )
+
+    hparams["datamodule"] = config["datamodule"]
+    hparams["trainer"] = config["trainer"]
+
+    if "seed" in config:
+        hparams["seed"] = config["seed"]
+    if "callbacks" in config:
+        hparams["callbacks"] = config["callbacks"]
+
+    logger.experiment.config.update(hparams)