added AST model

TODO.md

@@ -1,4 +1,5 @@
 - ✅ Ensure app.py audio input sounds like training data
+- ✅ Use a huggingface transformer with the dataset
 - Verify that the training spectrogram matches the predict spectrogram
 - Count number of example misses in dataset loading
 - Verify windowing and jitter params in Song Dataset
@@ -7,4 +8,4 @@
 - Verify that labels really match what is on the music4dance site
 - Read the Medium series about audio DL
 - double check \_rectify_duration
-- Filter out songs that have only one vote
+- ✅ Filter out songs that have only one vote

models/audio_spectrogram_transformer.py

@@ -0,0 +1,72 @@
+from transformers import ASTFeatureExtractor, AutoFeatureExtractor, ASTConfig, AutoModelForAudioClassification, TrainingArguments, Trainer
+import torch
+from torch import nn
+from sklearn.utils.class_weight import compute_class_weight
+import evaluate
+import numpy as np
+
+accuracy = evaluate.load("accuracy")
+
+def compute_metrics(eval_pred):
+    predictions = np.argmax(eval_pred.predictions, axis=1)
+    return accuracy.compute(predictions=predictions, references=eval_pred.label_ids)
+
+def get_id_label_mapping(labels:list[str]) -> tuple[dict, dict]:
+    id2label = {str(i) : label for i, label in enumerate(labels)}
+    label2id = {label : str(i) for i, label in enumerate(labels)}
+
+    return id2label, label2id
+
+def train(
+        labels,
+        train_ds, 
+        test_ds, 
+        output_dir="models/weights/ast",
+        device="cpu",
+        batch_size=128,
+        epochs=10):
+    id2label, label2id = get_id_label_mapping(labels)
+    model_checkpoint = "MIT/ast-finetuned-audioset-10-10-0.4593"
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_checkpoint)
+    preprocess_waveform = lambda wf : feature_extractor(wf, sampling_rate=train_ds.resample_frequency, padding="max_length", return_tensors="pt")
+    train_ds.map(preprocess_waveform)
+    test_ds.map(preprocess_waveform)
+
+    model = AutoModelForAudioClassification.from_pretrained(
+    model_checkpoint, 
+    num_labels=len(labels),
+    label2id=label2id,
+    id2label=id2label,
+    ignore_mismatched_sizes=True
+).to(device)
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        learning_rate=5e-5,
+        per_device_train_batch_size=batch_size,
+        gradient_accumulation_steps=5,
+        per_device_eval_batch_size=batch_size,
+        num_train_epochs=epochs,
+        warmup_ratio=0.1,
+        logging_steps=10,
+        load_best_model_at_end=True,
+        metric_for_best_model="accuracy",
+        push_to_hub=False,
+        use_mps_device=device == "mps"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=test_ds,
+        tokenizer=feature_extractor,
+        compute_metrics=compute_metrics,
+    )
+    trainer.train()
+    return model
+
+
+    
+

models/config/train.yaml

@@ -1,4 +1,5 @@
 global:
+  id: ast_ptl
   device: mps
   seed: 42
   dance_ids:
@@ -18,11 +19,11 @@ global:
     - VWZ
     - WCS
 data_module:
-  batch_size: 1024
-  num_workers: 10
-  min_votes: 2
-  song_data_path: data/songs_cleaned.csv
+  song_data_path: data/samples/songs_cleaned.csv
   song_audio_path: data/samples
+  batch_size: 256
+  num_workers: 10
+  min_votes: 1
   dataset_kwargs:
     audio_window_duration: 6
     audio_window_jitter: 1.5
@@ -40,7 +41,8 @@ trainer:
   fast_dev_run: False
   track_grad_norm: 2
   # gradient_clip_val: 0.5
+  overfit_batches: 1
 training_environment:
-  learning_rate: 0.0033
+  learning_rate: 0.00053
 model:
   n_channels: 128

models/residual.py

@@ -136,6 +136,7 @@ class TrainingEnvironment(pl.LightningModule):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
         # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min') {"scheduler": scheduler, "monitor": "val/loss"}
         return [optimizer] 
+    
         
 
 class DancePredictor:

preprocessing/dataset.py

@@ -7,6 +7,9 @@ from .pipelines import AudioTrainingPipeline
 import pytorch_lightning as pl
 from .preprocess import get_examples
 from sklearn.model_selection import train_test_split
+from torchaudio import transforms as taT
+from torch import nn
+from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
 
 
 
@@ -81,6 +84,54 @@ class SongDataset(Dataset):
 
     def _label_from_index(self, idx:int) -> torch.Tensor:
         return torch.from_numpy(self.dance_labels[self._convert_idx(idx)])
+    
+
+class WaveformSongDataset(SongDataset):
+    """
+    Outputs raw waveforms of the data instead of a spectrogram.
+    """
+
+    def __init__(self, *args,resample_frequency=16000, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.resample_frequency = resample_frequency
+        self.resampler = taT.Resample(self.sample_rate, self.resample_frequency)
+        self.pipeline = []
+
+    def __getitem__(self, idx:int) -> dict[str, torch.Tensor]:
+        waveform = self._waveform_from_index(idx)
+        assert waveform.shape[1] > 10, f"No data found: {self._backtrace_audio_path(idx)}"
+        # resample the waveform
+        waveform = self.resampler(waveform)
+        
+        waveform = waveform.mean(0)
+
+        dance_labels = self._label_from_index(idx)
+        return waveform, dance_labels
+    
+    
+
+
+class HuggingFaceWaveformSongDataset(WaveformSongDataset):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pipeline = []
+
+
+    def __getitem__(self, idx:int) -> dict[str, torch.Tensor]:
+        x,y = super().__getitem__(idx)
+        if len(self.pipeline) > 0:
+            for fn in self.pipeline:
+                x = fn(x)
+
+        dance_labels = y.argmax()
+        return {"input_values": x["input_values"][0] if hasattr(x, "input_values") else x, "label": dance_labels}
+
+    def map(self,fn):
+        """
+        NOTE this mutates the original, doesn't return a copy like normal maps.
+        """
+        self.pipeline.append(fn)
 
 class DanceDataModule(pl.LightningDataModule):
     def __init__(self, 
@@ -92,6 +143,7 @@ class DanceDataModule(pl.LightningDataModule):
     min_votes=1,
     batch_size:int=64,
     num_workers=10,
+    dataset_cls = None,
     dataset_kwargs={}
     ):
         super().__init__()
@@ -104,6 +156,7 @@ class DanceDataModule(pl.LightningDataModule):
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.dataset_kwargs = dataset_kwargs
+        self.dataset_cls = dataset_cls if dataset_cls is not None else SongDataset
 
         df = pd.read_csv(song_data_path)
         self.x,self.y = get_examples(df, self.song_audio_path,class_list=self.target_classes, multi_label=True, min_votes=min_votes)
@@ -115,7 +168,7 @@ class DanceDataModule(pl.LightningDataModule):
         self.test_ds = self._dataset_from_indices(test_i)
     
     def _dataset_from_indices(self, idx:list[int]) -> SongDataset:
-        return SongDataset(self.x[idx], self.y[idx], **self.dataset_kwargs)
+        return self.dataset_cls(self.x[idx], self.y[idx], **self.dataset_kwargs)
         
     def train_dataloader(self):
         return DataLoader(self.train_ds, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True)
@@ -128,4 +181,78 @@ class DanceDataModule(pl.LightningDataModule):
 
     def get_label_weights(self):
         n_examples, n_classes = self.y.shape
-        return torch.from_numpy(n_examples / (n_classes * sum(self.y)))
+        return torch.from_numpy(n_examples / (n_classes * sum(self.y)))
+    
+
+class WaveformTrainingEnvironment(pl.LightningModule):
+
+    def __init__(self, model: nn.Module, criterion: nn.Module, feature_extractor, config:dict, learning_rate=1e-4, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model = model
+        self.criterion = criterion
+        self.learning_rate = learning_rate
+        self.config=config
+        self.feature_extractor=feature_extractor
+        self.save_hyperparameters({
+            "model": type(model).__name__,
+            "loss": type(criterion).__name__,
+            "config": config,
+             **kwargs
+            })
+
+    def preprocess_inputs(self, x):
+        device = x.device
+        x = x.squeeze(1).cpu().numpy()
+        x = self.feature_extractor(list(x),return_tensors='pt', sampling_rate=16000)
+        return x["input_values"].to(device)
+    
+    def training_step(self, batch: tuple[torch.Tensor, torch.TensorType], batch_index: int) -> torch.Tensor:
+        features, labels = batch
+        features = self.preprocess_inputs(features)
+        outputs = self.model(features).logits
+        outputs = nn.Sigmoid()(outputs) # good for multi label classification, should be softmax otherwise
+        loss = self.criterion(outputs, labels)
+        metrics = calculate_metrics(outputs, labels, prefix="train/", multi_label=True)
+        self.log_dict(metrics, prog_bar=True)
+        return loss
+
+
+    def validation_step(self, batch:tuple[torch.Tensor, torch.TensorType], batch_index:int):
+        x,y = batch
+        x = self.preprocess_inputs(x)
+        preds = self.model(x).logits
+        preds = nn.Sigmoid()(preds) 
+        metrics = calculate_metrics(preds, y, prefix="val/", multi_label=True)
+        metrics["val/loss"] = self.criterion(preds, y)
+        self.log_dict(metrics,prog_bar=True)
+
+    def test_step(self, batch:tuple[torch.Tensor, torch.TensorType], batch_index:int):
+        x, y = batch
+        x = self.preprocess_inputs(x)
+        preds = self.model(x).logits
+        preds = nn.Sigmoid()(preds) 
+        self.log_dict(calculate_metrics(preds, y, prefix="test/", multi_label=True), prog_bar=True)
+    
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min') {"scheduler": scheduler, "monitor": "val/loss"}
+        return [optimizer] 
+    
+
+
+def calculate_metrics(pred, target, threshold=0.5, prefix="", multi_label=True) -> dict[str, torch.Tensor]:
+    target = target.detach().cpu().numpy()
+    pred = pred.detach().cpu().numpy()
+    params = {
+            "y_true": target if multi_label else target.argmax(1) ,
+            "y_pred": np.array(pred > threshold, dtype=float) if multi_label else pred.argmax(1), 
+            "zero_division": 0,
+            "average":"macro"
+            }
+    metrics= {
+            'precision': precision_score(**params),
+            'recall': recall_score(**params),
+            'f1': f1_score(**params),
+            'accuracy': accuracy_score(y_true=params["y_true"], y_pred=params["y_pred"]),
+            }
+    return {prefix + k: torch.tensor(v,dtype=torch.float32) for k,v in metrics.items()}

preprocessing/preprocess.py

@@ -11,12 +11,11 @@ from tqdm import tqdm
 def url_to_filename(url:str) -> str:
     return f"{url.split('/')[-1]}.wav"
 
-def get_songs_with_audio(df:pd.DataFrame, audio_dir:str) -> pd.DataFrame:
-    audio_urls = df["Sample"].replace(".", np.nan)
+def has_valid_audio(audio_urls:pd.Series, audio_dir:str) -> pd.Series:
+    audio_urls = audio_urls.replace(".", np.nan)
     audio_files = set(os.path.basename(f) for f in Path(audio_dir).iterdir())
-    valid_audio = audio_urls.apply(lambda url : url is not np.nan and url_to_filename(url) in audio_files)
-    df = df[valid_audio]
-    return df
+    valid_audio_mask = audio_urls.apply(lambda url : url is not np.nan and url_to_filename(url) in audio_files)
+    return valid_audio_mask
 
 def validate_audio(audio_urls:pd.Series, audio_dir:str) -> pd.Series:
     """
@@ -95,11 +94,11 @@ def vectorize_multi_label(labels: dict[str,int], unique_labels:np.ndarray) -> np
     return probs
 
 def get_examples(df:pd.DataFrame, audio_dir:str, class_list=None, multi_label=True, min_votes=1) -> tuple[np.ndarray, np.ndarray]:
-    sampled_songs = get_songs_with_audio(df, audio_dir)
-    sampled_songs.loc[:,"DanceRating"] = fix_dance_rating_counts(sampled_songs["DanceRating"])
+    sampled_songs = df[has_valid_audio(df["Sample"], audio_dir)]
+    sampled_songs["DanceRating"] = fix_dance_rating_counts(sampled_songs["DanceRating"])
     if class_list is not None:
         class_list = set(class_list)
-        sampled_songs.loc[:,"DanceRating"] = sampled_songs["DanceRating"].apply(
+        sampled_songs["DanceRating"] = sampled_songs["DanceRating"].apply(
             lambda labels : {k: v for k,v in labels.items() if k in class_list} 
             if not pd.isna(labels) and any(label in class_list and amt > 0 for label, amt in labels.items()) 
             else np.nan)

train.py

@@ -1,23 +1,47 @@
 from torch.utils.data import DataLoader
 import pandas as pd
+from typing import Callable
 from torch import nn
 from torch.utils.data import SubsetRandomSampler
 from sklearn.model_selection import KFold
 import pytorch_lightning as pl
 from pytorch_lightning import callbacks as cb
 from models.utils import LabelWeightedBCELoss
-from preprocessing.dataset import SongDataset
+from models.audio_spectrogram_transformer import train as train_audio_spectrogram_transformer, get_id_label_mapping
+from preprocessing.dataset import SongDataset, WaveformTrainingEnvironment
 from preprocessing.preprocess import get_examples
 from models.residual import ResidualDancer, TrainingEnvironment
 import yaml
-from preprocessing.dataset import DanceDataModule
+from preprocessing.dataset import DanceDataModule, WaveformSongDataset, HuggingFaceWaveformSongDataset
+from torch.utils.data import random_split
 from wakepy import keepawake
+import numpy as np
+from transformers import ASTFeatureExtractor, AutoFeatureExtractor, ASTConfig, AutoModelForAudioClassification
+from argparse import ArgumentParser
+
+
+
+import torch
+from torch import nn
+from sklearn.utils.class_weight import compute_class_weight
+
+def get_training_fn(id:str) -> Callable:
+    match id:
+        case "ast_ptl":
+            return train_ast_lightning
+        case "ast_hf":
+            return train_ast
+        case "residual_dancer":
+            return train_model
+        case _:
+            raise Exception(f"Couldn't find a training function for '{id}'.")
 
 def get_config(filepath:str) -> dict:
     with open(filepath, "r") as f:
         config = yaml.safe_load(f)
     return config
 
+
 def cross_validation(config, k=5):
     df = pd.read_csv("data/songs.csv")
     g_config = config["global"]
@@ -52,7 +76,8 @@ def train_model(config:dict):
         # cb.LearningRateFinder(update_attr=True),
         cb.EarlyStopping("val/loss", patience=5),
         cb.StochasticWeightAveraging(1e-2),
-        cb.RichProgressBar()
+        cb.RichProgressBar(),
+        cb.DeviceStatsMonitor(),
     ]
     trainer = pl.Trainer(
         callbacks=callbacks, 
@@ -62,8 +87,69 @@ def train_model(config:dict):
     trainer.test(train_env, datamodule=data)
 
 
+def train_ast(
+    config:dict
+):
+    TARGET_CLASSES = config["global"]["dance_ids"]
+    DEVICE = config["global"]["device"]
+    SEED = config["global"]["seed"]
+    dataset_kwargs = config["data_module"]["dataset_kwargs"]
+    test_proportion = config["data_module"].get("test_proportion", 0.2)
+    train_proportion = 1. - test_proportion
+    song_data_path="data/songs_cleaned.csv"
+    song_audio_path = "data/samples"
+    pl.seed_everything(SEED, workers=True)
+
+    df = pd.read_csv(song_data_path)
+    x, y = get_examples(df, song_audio_path,class_list=TARGET_CLASSES, multi_label=True)
+    train_i, test_i = random_split(np.arange(len(x)), [train_proportion, test_proportion])
+    train_ds = HuggingFaceWaveformSongDataset(x[train_i], y[train_i], **dataset_kwargs, resample_frequency=16000)
+    test_ds = HuggingFaceWaveformSongDataset(x[test_i], y[test_i], **dataset_kwargs, resample_frequency=16000)
+    train_audio_spectrogram_transformer(TARGET_CLASSES, train_ds, test_ds, device=DEVICE)
+
+
+def train_ast_lightning(config:dict):
+    """
+    work on integration between waveform dataset and environment. Should work for both HF and PTL.
+    """
+    TARGET_CLASSES = config["global"]["dance_ids"]
+    DEVICE = config["global"]["device"]
+    SEED = config["global"]["seed"]
+    pl.seed_everything(SEED, workers=True)
+    data = DanceDataModule(target_classes=TARGET_CLASSES, dataset_cls=WaveformSongDataset, **config['data_module'])
+    id2label, label2id = get_id_label_mapping(TARGET_CLASSES)
+    model_checkpoint = "MIT/ast-finetuned-audioset-10-10-0.4593"
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_checkpoint)
+
+    model = AutoModelForAudioClassification.from_pretrained(
+    model_checkpoint, 
+    num_labels=len(label2id),
+    label2id=label2id,
+    id2label=id2label,
+    ignore_mismatched_sizes=True
+).to(DEVICE)
+    label_weights = data.get_label_weights().to(DEVICE)
+    criterion = LabelWeightedBCELoss(label_weights) #nn.CrossEntropyLoss(label_weights)
+    train_env = WaveformTrainingEnvironment(model, criterion,feature_extractor, config)
+    callbacks = [
+        # cb.LearningRateFinder(update_attr=True),
+        cb.EarlyStopping("val/loss", patience=5),
+        cb.StochasticWeightAveraging(1e-2),
+        cb.RichProgressBar()
+    ]
+    trainer = pl.Trainer(
+        callbacks=callbacks, 
+        **config["trainer"]
+        )
+    trainer.fit(train_env, datamodule=data)
+    trainer.test(train_env, datamodule=data)
 
 if __name__ == "__main__":
-    config = get_config("models/config/train.yaml")
+    parser = ArgumentParser(description="Trains models on the dance dataset and saves weights.")
+    parser.add_argument("--config", help="Path to the yaml file that defines the training configuration.", default="models/config/train.yaml")
+    args = parser.parse_args()
+    config = get_config(args.config)
+    training_id = config["global"]["id"]
+    train = get_training_fn(training_id)
     with keepawake():
-        train_model(config)
+        train(config)