Merge pull request #6 from mhrice/metrics-init

README.md

@@ -3,16 +3,17 @@
 1. `python3 -m venv env`
 2. `source env/bin/activate`
 3. `pip install -e .`
-4. `pip install -e umx`
+4. `git submodule update --init --recursive`
+5. `pip install -e umx`
 
 ## Download [GuitarFX Dataset] (https://zenodo.org/record/7044411/)
-`./download_egfx.sh`
+`./scripts/download_egfx.sh`
 
 ## Train model
 1. Change Wandb variables in `shell_vars.sh`
-2. `python train.py exp=audio_diffusion`
+2. `python scripts/train.py exp=audio_diffusion`
 or
-2. `python train.py exp=umx`
+2. `python scripts/train.py exp=umx`
 
 To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-line
 

config.yaml

@@ -20,16 +20,18 @@ callbacks:
     filename: '{epoch:02d}-{valid_loss:.3f}'
 
 datamodule:
-  _target_: datasets.Datamodule
+  _target_: remfx.datasets.Datamodule
   dataset:
-    _target_: datasets.GuitarFXDataset
+    _target_: remfx.datasets.GuitarFXDataset
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     length: ${length}
+    chunk_size_in_sec: 6
   val_split: 0.2
   batch_size: 16
   num_workers: 8
   pin_memory: True
+  persistent_workers: True
 
 logger:
   _target_: pytorch_lightning.loggers.WandbLogger

remfx/datasets.py

@@ -1,10 +1,11 @@
+import torch
 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
 import pytorch_lightning as pl
-from typing import Any, List
+from typing import Any, List, Tuple
 
 # https://zenodo.org/record/7044411/
 
@@ -18,52 +19,98 @@ class GuitarFXDataset(Dataset):
         root: str,
         sample_rate: int,
         length: int = LENGTH,
+        chunk_size_in_sec: int = 3,
         effect_types: List[str] = None,
     ):
         self.length = length
         self.wet_files = []
         self.dry_files = []
+        self.chunks = []
         self.labels = []
+        self.song_idx = []
         self.root = Path(root)
+        self.chunk_size_in_sec = chunk_size_in_sec
 
         if effect_types is None:
             effect_types = [
                 d.name for d in self.root.iterdir() if d.is_dir() and d != "Clean"
             ]
+        current_file = 0
         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(
+                wet_files = sorted(list(pickup.glob("*.wav")))
+                dry_files = sorted(
                     list(self.root.glob(f"Clean/{pickup.name}/**/*.wav"))
                 )
-                self.labels += [i] * len(self.wet_files)
+                self.wet_files += wet_files
+                self.dry_files += dry_files
+                self.labels += [i] * len(wet_files)
+                for audio_file in wet_files:
+                    chunk_starts = create_sequential_chunks(
+                        audio_file, self.chunk_size_in_sec
+                    )
+                    self.chunks += chunk_starts
+                    self.song_idx += [current_file] * len(chunk_starts)
+                    current_file += 1
         print(
-            f"Found {len(self.wet_files)} wet files and {len(self.dry_files)} dry files"
+            f"Found {len(self.wet_files)} wet files and {len(self.dry_files)} dry files.\n"
+            f"Total chunks: {len(self.chunks)}"
         )
         self.resampler = T.Resample(ORIG_SR, sample_rate)
 
     def __len__(self):
-        return len(self.dry_files)
+        return len(self.chunks)
 
     def __getitem__(self, idx):
-        x, sr = torchaudio.load(self.wet_files[idx])
-        y, sr = torchaudio.load(self.dry_files[idx])
-        effect_label = self.labels[idx]
+        # Load effected and "clean" audio
+        song_idx = self.song_idx[idx]
+        x, sr = torchaudio.load(self.wet_files[song_idx])
+        y, sr = torchaudio.load(self.dry_files[song_idx])
+        effect_label = self.labels[song_idx]  # Effect label
+
+        chunk_start = self.chunks[idx]
+        chunk_size_in_samples = self.chunk_size_in_sec * sr
+        x = x[:, chunk_start : chunk_start + chunk_size_in_samples]
+        y = y[:, chunk_start : chunk_start + chunk_size_in_samples]
 
         resampled_x = self.resampler(x)
         resampled_y = self.resampler(y)
-        # Pad or crop to length
+        # Pad to length if needed
         if resampled_x.shape[-1] < self.length:
             resampled_x = F.pad(resampled_x, (0, self.length - resampled_x.shape[1]))
-        elif resampled_x.shape[-1] > self.length:
-            resampled_x = resampled_x[:, : self.length]
         if resampled_y.shape[-1] < self.length:
             resampled_y = F.pad(resampled_y, (0, self.length - resampled_y.shape[1]))
-        elif resampled_y.shape[-1] > self.length:
-            resampled_y = resampled_y[:, : self.length]
         return (resampled_x, resampled_y, effect_label)
 
 
+def create_random_chunks(
+    audio_file: str, chunk_size: int, num_chunks: int
+) -> List[Tuple[int, int]]:
+    """Create num_chunks random chunks of size chunk_size (seconds)
+    from an audio file.
+    Return sample_index of start of each chunk
+    """
+    audio, sr = torchaudio.load(audio_file)
+    chunk_size_in_samples = chunk_size * sr
+    if chunk_size_in_samples >= audio.shape[-1]:
+        chunk_size_in_samples = audio.shape[-1] - 1
+    chunks = []
+    for i in range(num_chunks):
+        start = torch.randint(0, audio.shape[-1] - chunk_size_in_samples, (1,)).item()
+        chunks.append(start)
+    return chunks
+
+
+def create_sequential_chunks(audio_file: str, chunk_size: int) -> List[Tuple[int, int]]:
+    """Create sequential chunks of size chunk_size (seconds) from an audio file.
+    Return sample_index of start of each chunk
+    """
+    audio, sr = torchaudio.load(audio_file)
+    chunk_size_in_samples = chunk_size * sr
+    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size_in_samples)
+    return chunk_starts
+
+
 class Datamodule(pl.LightningDataModule):
     def __init__(
         self,

remfx/models.py

@@ -4,7 +4,9 @@ import pytorch_lightning as pl
 from einops import rearrange
 import wandb
 from audio_diffusion_pytorch import DiffusionModel
-import auraloss
+from auraloss.time import SISDRLoss
+from auraloss.freq import MultiResolutionSTFTLoss, STFTLoss
+from torch.nn import L1Loss
 
 from umx.openunmix.model import OpenUnmix, Separator
 
@@ -28,6 +30,13 @@ class RemFXModel(pl.LightningModule):
         self.lr_weight_decay = lr_weight_decay
         self.sample_rate = sample_rate
         self.model = network
+        self.metrics = torch.nn.ModuleDict(
+            {
+                "SISDR": SISDRLoss(),
+                "STFT": STFTLoss(),
+                "L1": L1Loss(),
+            }
+        )
 
     @property
     def device(self):
@@ -49,10 +58,24 @@ class RemFXModel(pl.LightningModule):
 
     def validation_step(self, batch, batch_idx):
         loss = self.common_step(batch, batch_idx, mode="valid")
+        return loss
 
     def common_step(self, batch, batch_idx, mode: str = "train"):
-        loss = self.model(batch)
+        loss, output = self.model(batch)
         self.log(f"{mode}_loss", loss)
+        x, y, label = batch
+        # Metric logging
+        for metric in self.metrics:
+            self.log(
+                f"{mode}_{metric}",
+                self.metrics[metric](output, y),
+                on_step=False,
+                on_epoch=True,
+                logger=True,
+                prog_bar=True,
+                sync_dist=True,
+            )
+
         return loss
 
     def on_validation_epoch_start(self):
@@ -61,29 +84,21 @@ class RemFXModel(pl.LightningModule):
     def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
         if self.log_next:
             x, target, label = batch
-            y = self.model.sample(x)
-            log_wandb_audio_batch(
-                logger=self.logger,
-                id="sample",
-                samples=x.cpu(),
-                sampling_rate=self.sample_rate,
-                caption=f"Epoch {self.current_epoch}",
-            )
-            log_wandb_audio_batch(
-                logger=self.logger,
-                id="prediction",
-                samples=y.cpu(),
-                sampling_rate=self.sample_rate,
-                caption=f"Epoch {self.current_epoch}",
-            )
+            self.model.eval()
+            with torch.no_grad():
+                y = self.model.sample(x)
+
+            # Concat samples together for easier viewing in dashboard
+            concat_samples = torch.cat([y, x, target], dim=-1)
             log_wandb_audio_batch(
                 logger=self.logger,
-                id="target",
-                samples=target.cpu(),
+                id="prediction_input_target",
+                samples=concat_samples.cpu(),
                 sampling_rate=self.sample_rate,
                 caption=f"Epoch {self.current_epoch}",
             )
             self.log_next = False
+            self.model.train()
 
 
 class OpenUnmixModel(torch.nn.Module):
@@ -116,7 +131,7 @@ class OpenUnmixModel(torch.nn.Module):
             n_fft=self.n_fft,
             n_hop=self.hop_length,
         )
-        self.loss_fn = auraloss.freq.MultiResolutionSTFTLoss(
+        self.loss_fn = MultiResolutionSTFTLoss(
             n_bins=self.num_bins, sample_rate=self.sample_rate
         )
 
@@ -127,7 +142,7 @@ class OpenUnmixModel(torch.nn.Module):
         sep_out = self.separator(x).squeeze(1)
         loss = self.loss_fn(sep_out, target)
 
-        return loss
+        return loss, sep_out
 
     def sample(self, x: Tensor) -> Tensor:
         return self.separator(x).squeeze(1)
@@ -140,7 +155,8 @@ class DiffusionGenerationModel(nn.Module):
 
     def forward(self, batch):
         x, target, label = batch
-        return self.model(x)
+        sampled_out = self.model.sample(x)
+        return self.model(x), sampled_out
 
     def sample(self, x: Tensor, num_steps: int = 10) -> Tensor:
         noise = torch.randn(x.shape).to(x)

scripts/download_egfx.sh

@@ -17,5 +17,6 @@ wget https://zenodo.org/record/7044411/files/Sweep-Echo.zip?download=1 -O Sweep-
 wget https://zenodo.org/record/7044411/files/TapeEcho.zip?download=1 -O TapeEcho.zip
 wget https://zenodo.org/record/7044411/files/TubeScreamer.zip?download=1 -O TubeScreamer.zip
 unzip -n \*.zip
+rm -rf *.zip
 
 

scripts/train.py

@@ -6,15 +6,14 @@ import remfx.utils as utils
 log = utils.get_logger(__name__)
 
 
-@hydra.main(version_base=None, config_path=".", config_name="config.yaml")
+@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)
+    if cfg.seed:
+        pl.seed_everything(cfg.seed)
 
     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")