refactored loggers

TODO.md

@@ -12,6 +12,8 @@
 - ✅ Download songs from [Best Ballroom](https://www.youtube.com/channel/UC0bYSnzAFMwPiEjmVsrvmRg)
 
 - ✅ fix nan values
+- Try higher mels (224) and more ffts (2048)
+- Verify random sample of dataset outputs by hand.
 
 ## Notes
 

audio_utils.py

@@ -1,42 +0,0 @@
-import librosa
-from IPython.display import Audio, display
-import matplotlib.pyplot as plt
-import torch
-SAMPLE_RIR_URL = "https://pytorch-tutorial-assets.s3.amazonaws.com/VOiCES_devkit/distant-16k/room-response/rm1/impulse/Lab41-SRI-VOiCES-rm1-impulse-mc01-stu-clo.wav"
-
-SAMPLE_NOISE_URL = "https://pytorch-tutorial-assets.s3.amazonaws.com/VOiCES_devkit/distant-16k/distractors/rm1/babb/Lab41-SRI-VOiCES-rm1-babb-mc01-stu-clo.wav"
-
-def plot_spectrogram(spec, title=None, ylabel="freq_bin", aspect="auto", xmax=None):
-    spec = spec.squeeze(0)
-    spec = spec.numpy()
-    fig, axs = plt.subplots(1, 1)
-    axs.set_title(title or "Spectrogram (db)")
-    axs.set_ylabel(ylabel)
-    axs.set_xlabel("frame")
-    im = axs.imshow(librosa.power_to_db(spec), origin="lower", aspect=aspect)
-    if xmax:
-        axs.set_xlim((0, xmax))
-    fig.colorbar(im, ax=axs)
-    plt.show(block=False)
-
-def play_audio(waveform, sample_rate):
-    waveform = waveform.numpy()
-
-    num_channels, num_frames = waveform.shape
-    if num_channels == 1:
-        display(Audio(waveform[0], rate=sample_rate))
-    elif num_channels == 2:
-        display(Audio((waveform[0], waveform[1]), rate=sample_rate))
-    else:
-        raise ValueError("Waveform with more than 2 channels are not supported.")
-
-def get_rir_sample(path, resample=None, processed=False):
-    rir_raw, sample_rate = torch.load(path)
-    if not processed:
-        return rir_raw, sample_rate
-    rir = rir_raw[:, int(sample_rate*1.01):int(sample_rate*1.3)]
-    rir = rir / torch.norm(rir, p=2)
-    rir = torch.flip(rir, [1])
-    return rir, sample_rate
-
-

models/config/train_local.yaml

@@ -19,6 +19,7 @@ dance_ids: &dance_ids
 data_module:
   batch_size: 64
   num_workers: 10
+  data_subset: 0.1
   test_proportion: 0.2
 
 datasets:
@@ -55,4 +56,6 @@ trainer:
 
 training_environment:
   learning_rate: 0.00053
-  log_spectrograms: False
+  loggers:
+    models.training_environment.SpectrogramLogger:
+      frequency: 100

models/training_environment.py

@@ -1,6 +1,7 @@
+import importlib
 from models.utils import calculate_metrics
 
-
+from abc import ABC, abstractmethod
 import pytorch_lightning as pl
 import torch
 import torch.nn as nn
@@ -13,7 +14,6 @@ class TrainingEnvironment(pl.LightningModule):
         criterion: nn.Module,
         config: dict,
         learning_rate=1e-4,
-        log_spectrograms=False,
         *args,
         **kwargs,
     ):
@@ -21,7 +21,9 @@ class TrainingEnvironment(pl.LightningModule):
         self.model = model
         self.criterion = criterion
         self.learning_rate = learning_rate
-        self.log_spectrograms = log_spectrograms
+        self.experiment_loggers = load_loggers(
+            config["training_environment"].get("loggers", {})
+        )
         self.config = config
         self.has_multi_label_predictions = (
             not type(criterion).__name__ == "CrossEntropyLoss"
@@ -48,15 +50,9 @@ class TrainingEnvironment(pl.LightningModule):
             multi_label=self.has_multi_label_predictions,
         )
         self.log_dict(metrics, prog_bar=True)
-        # Log spectrograms
-        if self.log_spectrograms and batch_index % 100 == 0:
-            tensorboard = self.logger.experiment
-            img_index = torch.randint(0, len(features), (1,)).item()
-            img = features[img_index][0]
-            img = (img - img.min()) / (img.max() - img.min())
-            tensorboard.add_image(
-                f"batch: {batch_index}, element: {img_index}", img, 0, dataformats="HW"
-            )
+        experiment = self.logger.experiment
+        for logger in self.experiment_loggers:
+            logger.step(experiment, batch_index, features, labels)
         return loss
 
     def validation_step(
@@ -88,3 +84,36 @@ class TrainingEnvironment(pl.LightningModule):
             "lr_scheduler": scheduler,
             "monitor": "val/loss",
         }
+
+
+class ExperimentLogger(ABC):
+    @abstractmethod
+    def step(self, experiment, data):
+        pass
+
+
+class SpectrogramLogger(ExperimentLogger):
+    def __init__(self, frequency=100) -> None:
+        self.frequency = frequency
+        self.counter = 0
+
+    def step(self, experiment, batch_index, x, label):
+        if self.counter == self.frequency:
+            self.counter = 0
+            img_index = torch.randint(0, len(x), (1,)).item()
+            img = x[img_index][0]
+            img = (img - img.min()) / (img.max() - img.min())
+            experiment.add_image(
+                f"batch: {batch_index}, element: {img_index}", img, 0, dataformats="HW"
+            )
+        self.counter += 1
+
+
+def load_loggers(logger_config: dict) -> list[ExperimentLogger]:
+    loggers = []
+    for logger_path, kwargs in logger_config.items():
+        module_name, class_name = logger_path.rsplit(".", 1)
+        module = importlib.import_module(module_name)
+        Logger = getattr(module, class_name)
+        loggers.append(Logger(**kwargs))
+    return loggers