fixed weighing strategy

environment.yml

@@ -1,4 +1,4 @@
-name: dancer-net
+name: dance-classifier
 channels:
   - pytorch
   - nvidia

models/audio_spectrogram_transformer.py

@@ -20,8 +20,7 @@ from preprocessing.dataset import (
     HuggingFaceDatasetWrapper,
     get_datasets,
 )
-from preprocessing.dataset import get_music4dance_examples
-from .utils import get_id_label_mapping, compute_hf_metrics
+from .utils import LabelWeightedBCELoss, get_id_label_mapping, compute_hf_metrics
 
 import pytorch_lightning as pl
 from pytorch_lightning import callbacks as cb
@@ -54,7 +53,7 @@ class AST(nn.Module):
 class ASTExtractorWrapper:
     def __init__(self, sampling_rate=16000, return_tensors="pt") -> None:
         max_length = 1024
-        self.extractor = ASTFeatureExtractor(do_normalize=False, max_length=max_length)
+        self.extractor = ASTFeatureExtractor(max_length=max_length, do_normalize=True)
         self.sampling_rate = sampling_rate
         self.return_tensors = return_tensors
         self.waveform_pipeline = WaveformTrainingPipeline()  # TODO configure from yaml
@@ -68,8 +67,6 @@ class ASTExtractorWrapper:
         )
 
         x = x["input_values"].squeeze(0).to(device)
-        # normalize
-        x = (x - x.mean()) / x.std()
         return x
 
 
@@ -90,9 +87,7 @@ def train_lightning_ast(config: dict):
     )
     model = AST(TARGET_CLASSES).to(DEVICE)
     label_weights = data.get_label_weights().to(DEVICE)
-    criterion = nn.CrossEntropyLoss(
-        label_weights
-    )  # LabelWeightedBCELoss(label_weights)
+    criterion = LabelWeightedBCELoss(label_weights)
     if "checkpoint" in config:
         train_env = TrainingEnvironment.load_from_checkpoint(
             config["checkpoint"], criterion=criterion, model=model, config=config
@@ -100,8 +95,7 @@ def train_lightning_ast(config: dict):
     else:
         train_env = TrainingEnvironment(model, criterion, config)
     callbacks = [
-        # cb.LearningRateFinder(update_attr=True),
-        cb.EarlyStopping("val/loss", patience=5),
+        cb.EarlyStopping("val/loss", patience=2),
         cb.RichProgressBar(),
     ]
     trainer = pl.Trainer(callbacks=callbacks, **config["trainer"])

models/config/train_local.yaml

@@ -4,7 +4,7 @@ seed: 42
 dance_ids: &dance_ids
   - BCH
   - BOL
-  - CHA
+  # - CHA
   - ECS
   - HST
   - LHP

models/training_environment.py

@@ -60,9 +60,6 @@ class TrainingEnvironment(pl.LightningModule):
             multi_label=self.has_multi_label_predictions,
         )
         self.log_dict(metrics, prog_bar=True)
-        experiment = self.logger.experiment
-        for logger in self.experiment_loggers:
-            logger.step(experiment, batch_index, features, labels)
         return loss
 
     def validation_step(
@@ -117,8 +114,10 @@ class TrainingEnvironment(pl.LightningModule):
         dance_ids = sorted(self.config["dance_ids"])
         np.fill_diagonal(self.test_cm, 0)
         cm = self.test_cm / self.test_cm.max()
-        ConfusionMatrixDisplay(cm, display_labels=dance_ids).plot()
-        image = plot_to_image(plt.gcf())
+        cm_plot = ConfusionMatrixDisplay(cm, display_labels=dance_ids)
+        fig, ax = plt.subplots(figsize=(12, 12))
+        cm_plot.plot(ax=ax)
+        image = plot_to_image(fig)
         image = torch.tensor(image, dtype=torch.uint8)
         image = image.permute(2, 0, 1)
         self.logger.experiment.add_image("test/confusion_matrix", image, 0)

preprocessing/dataset.py

@@ -99,10 +99,21 @@ class SongDataset(Dataset):
             total_slices += audio_slices
 
     def get_label_weights(self):
-        n_examples, n_classes = self.dance_labels.shape
-        weights = n_examples / (n_classes * sum(self.dance_labels))
-        weights[np.isinf(weights)] = 0.0
-        return torch.from_numpy(weights)
+        n_examples = len(self)
+        n_classes = self.dance_labels.shape[1]
+        dance_label_counts = self.get_dance_label_counts()
+        weights = n_examples / (n_classes * dance_label_counts)
+        weights[np.isinf(weights) | np.isnan(weights)] = 1.0
+        return torch.from_numpy(weights).type(torch.float32)
+
+    def get_dance_label_counts(self) -> np.ndarray:
+        """
+        Returns the number of examples for each dance label.
+        """
+        examples_per_audio = np.expand_dims(
+            np.array(self.audio_durations) // self.audio_window_duration, axis=-1
+        )
+        return sum(self.dance_labels * examples_per_audio)
 
     def _backtrace_audio_path(self, index: int) -> str:
         return self.audio_paths[self._idx2audio_idx(index)]
@@ -230,6 +241,7 @@ class Music4DanceDataset(Dataset):
         class_list=None,
         multi_label=True,
         min_votes=1,
+        class_count_limit=None,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -240,6 +252,7 @@ class Music4DanceDataset(Dataset):
             class_list=class_list,
             multi_label=multi_label,
             min_votes=min_votes,
+            class_count_limit=class_count_limit,
         )
         self.song_dataset = SongDataset(
             song_paths,
@@ -256,7 +269,12 @@ class Music4DanceDataset(Dataset):
 
 
 def get_music4dance_examples(
-    df: pd.DataFrame, audio_dir: str, class_list=None, multi_label=True, min_votes=1
+    df: pd.DataFrame,
+    audio_dir: str,
+    class_list=None,
+    multi_label=True,
+    min_votes=1,
+    class_count_limit=None,
 ) -> tuple[np.ndarray, np.ndarray]:
     sampled_songs = df[has_valid_audio(df["Sample"], audio_dir)].copy(deep=True)
     sampled_songs["DanceRating"] = fix_dance_rating_counts(sampled_songs["DanceRating"])

preprocessing/preprocess.py

@@ -10,9 +10,7 @@ import torchaudio
 import torch
 from tqdm import tqdm
 
-
-def url_to_filename(url: str) -> str:
-    return f"{url.split('/')[-1]}.wav"
+from preprocessing.utils import url_to_filename
 
 
 def has_valid_audio(audio_urls: pd.Series, audio_dir: str) -> pd.Series:

preprocessing/utils.py

@@ -0,0 +1,32 @@
+import sounddevice as sd
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def url_to_filename(url: str) -> str:
+    return f"{url.split('/')[-1]}.wav"
+
+
+def play_audio(waveform: np.ndarray, sample_rate: int):
+    """
+    Assumes that waveform is a numpy array normalized between -1 and 1.
+    """
+    if waveform.max() > 1.0 or waveform.min() < -1.0:
+        raise ValueError("waveform must be a numpy array normalized between -1 and 1.")
+    sd.play(waveform, sample_rate)
+    sd.wait()
+
+
+def plot_spectrogram(spec, title=None, ylabel="freq_bin", aspect="auto", xmax=None):
+    """
+    Assumes that the spectrogram is in decibels.
+    """
+    fig, axs = plt.subplots(1, 1)
+    axs.set_title(title or "Spectrogram (db)")
+    axs.set_ylabel(ylabel)
+    axs.set_xlabel("frame")
+    im = axs.imshow(spec, origin="lower", aspect=aspect)
+    if xmax:
+        axs.set_xlim((0, xmax))
+    fig.colorbar(im, ax=axs)
+    return fig