ex_pl_audioset.py

import wandb
import numpy as np
import os
from torch import autocast
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
import argparse
from sklearn import metrics
from contextlib import nullcontext
import torch.nn as nn
import torch.nn.functional as F
from torch.hub import download_url_to_file
import pickle
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import LearningRateMonitor

from datasets.audioset import get_test_set, get_full_training_set, get_ft_weighted_sampler
from models.mn.model import get_model as get_mobilenet
from models.dymn.model import get_model as get_dymn
from models.ensemble import get_ensemble_model
from models.preprocess import AugmentMelSTFT
from helpers.init import worker_init_fn
from helpers.utils import NAME_TO_WIDTH, exp_warmup_linear_down, mixup

preds_url = \
    "https://github.com/fschmid56/EfficientAT/releases/download/v0.0.1/passt_enemble_logits_mAP_495.npy"

fname_to_index_url = "https://github.com/fschmid56/EfficientAT/releases/download/v0.0.1/fname_to_index.pkl"


class PLModule(pl.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        # model to preprocess waveform to mel spectrograms
        self.mel = AugmentMelSTFT(n_mels=config.n_mels,
                         sr=config.resample_rate,
                         win_length=config.window_size,
                         hopsize=config.hop_size,
                         n_fft=config.n_fft,
                         freqm=config.freqm,
                         timem=config.timem,
                         fmin=config.fmin,
                         fmax=config.fmax,
                         fmin_aug_range=config.fmin_aug_range,
                         fmax_aug_range=config.fmax_aug_range
                         )

        # load prediction model
        model_name = config.model_name
        pretrained_name = model_name if config.pretrained else None
        width = NAME_TO_WIDTH(model_name) if model_name and config.pretrained else config.model_width
        if model_name.startswith("dymn"):
            model = get_dymn(width_mult=width, pretrained_name=pretrained_name,
                             strides=config.strides, pretrain_final_temp=config.pretrain_final_temp)
        else:
            model = get_mobilenet(width_mult=width, pretrained_name=pretrained_name,
                                  strides=config.strides, head_type=config.head_type, se_dims=config.se_dims)
        self.model = model

        # prepare ingredients for knowledge distillation
        assert 0 <= config.kd_lambda <= 1, "Lambda for Knowledge Distillation must be between 0 and 1."
        self.distillation_loss = nn.BCEWithLogitsLoss(reduction="none")

        # load stored teacher predictions
        if not os.path.isfile(config.teacher_preds):
            # download file
            print("Download teacher predictions...")
            download_url_to_file(preds_url, config.teacher_preds)
        print(f"Load teacher predictions from file {config.teacher_preds}")
        teacher_preds = np.load(config.teacher_preds)
        teacher_preds = torch.from_numpy(teacher_preds).float()
        teacher_preds = torch.sigmoid(teacher_preds / config.temperature)
        teacher_preds.requires_grad = False
        self.teacher_preds = teacher_preds

        if not os.path.isfile(config.fname_to_index):
            print("Download filename to teacher prediction index dictionary...")
            download_url_to_file(fname_to_index_url, config.fname_to_index)
        with open(config.fname_to_index, 'rb') as f:
            fname_to_index = pickle.load(f)
        self.fname_to_index = fname_to_index

        self.distributed_mode = config.num_devices > 1
        self.training_step_outputs = []
        self.validation_step_outputs = []

    def mel_forward(self, x):
        old_shape = x.size()
        x = x.reshape(-1, old_shape[2])
        x = self.mel(x)
        x = x.reshape(old_shape[0], old_shape[1], x.shape[1], x.shape[2])
        return x

    def forward(self, x):
        """
        :param x: batch of raw audio signals (waveforms)
        :return: final model predictions
        """
        x = self.mel_forward(x)
        x = self.model(x)
        return x

    def configure_optimizers(self):
        """
        This is the way pytorch lightening requires optimizers and learning rate schedulers to be defined.
        The specified items are used automatically in the optimization loop (no need to call optimizer.step() yourself).
        :return: dict containing optimizer and learning rate scheduler
        """
        if self.config.adamw:
            optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.config.max_lr,
                                          weight_decay=self.config.weight_decay)
        else:
            optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.max_lr,
                                         weight_decay=self.config.weight_decay)

        # phases of lr schedule: exponential increase, constant lr, linear decrease, fine-tune
        schedule_lambda = \
            exp_warmup_linear_down(self.config.warm_up_len, self.config.ramp_down_len, self.config.ramp_down_start,
                                   self.config.last_lr_value)
        lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule_lambda)

        return {
            'optimizer': optimizer,
            'lr_scheduler': lr_scheduler
        }

    def on_train_epoch_start(self):
        # in case of DyMN: update DyConv temperature
        if hasattr(self.model, "update_params"):
            self.model.update_params(self.current_epoch)

    def training_step(self, train_batch, batch_idx):
        """
        :param train_batch: contains one batch from train dataloader
        :param batch_idx
        :return: a dict containing at least loss that is used to update model parameters, can also contain
                    other items that can be processed in 'training_epoch_end' to log other metrics than loss
        """
        x, f, y, i = train_batch
        bs = x.size(0)
        x = self.mel_forward(x)

        rn_indices, lam = None, None
        if self.config.mixup_alpha:
            rn_indices, lam = mixup(bs, self.config.mixup_alpha)
            lam = lam.to(x.device)
            x = x * lam.reshape(bs, 1, 1, 1) + \
                x[rn_indices] * (1. - lam.reshape(bs, 1, 1, 1))
            y_hat, _ = self.model(x)
            y_mix = y * lam.reshape(bs, 1) + y[rn_indices] * (1. - lam.reshape(bs, 1))
            samples_loss = F.binary_cross_entropy_with_logits(y_hat, y_mix, reduction="none")
        else:
            y_hat, _ = self.model(x)
            samples_loss = F.binary_cross_entropy_with_logits(y_hat, y, reduction="none")

        # hard label loss
        label_loss = samples_loss.mean()

        # distillation loss
        if self.config.kd_lambda > 0:
            # fetch the correct index in 'teacher_preds' for given filename
            # insert -1 for files not in fname_to_index (proportion of files successfully downloaded from
            # YouTube can vary for AudioSet)
            indices = torch.tensor(
                [self.fname_to_index[fname] if fname in self.fname_to_index else -1 for fname in f], dtype=torch.int64
            )
            # get indices of files we could not find the teacher predictions for
            unknown_indices = indices == -1
            y_soft_teacher = self.teacher_preds[indices]
            y_soft_teacher = y_soft_teacher.to(y_hat.device).type_as(y_hat)

            if self.config.mixup_alpha:
                soft_targets_loss = \
                    self.distillation_loss(y_hat, y_soft_teacher).mean(dim=1) * lam.reshape(bs) + \
                    self.distillation_loss(y_hat, y_soft_teacher[rn_indices]).mean(dim=1) \
                    * (1. - lam.reshape(bs))
            else:
                soft_targets_loss = distillation_loss(y_hat, y_soft_teacher)

            # zero out loss for samples we don't have teacher predictions for
            soft_targets_loss[unknown_indices] = soft_targets_loss[unknown_indices] * 0
            soft_targets_loss = soft_targets_loss.mean()

            # weighting losses
            label_loss = self.config.kd_lambda * label_loss
            soft_targets_loss = (1 - self.config.kd_lambda) * soft_targets_loss
        else:
            soft_targets_loss = torch.tensor(0., device=label_loss.device, dtype=label_loss.dtype)

        # total loss is sum of lambda-weighted label and distillation loss
        loss = label_loss + soft_targets_loss

        results = {"loss": loss.detach().cpu(), "label_loss": label_loss.detach().cpu(),
                   "kd_loss": soft_targets_loss.detach().cpu()}
        self.training_step_outputs.append(results)
        return loss

    def on_train_epoch_end(self):
        """
        :return: a dict containing the metrics you want to log to Weights and Biases
        """
        avg_loss = torch.stack([x['loss'] for x in self.training_step_outputs]).mean()
        avg_label_loss = torch.stack([x['label_loss'] for x in self.training_step_outputs]).mean()
        avg_kd_loss = torch.stack([x['kd_loss'] for x in self.training_step_outputs]).mean()
        self.log_dict({'train/loss': torch.as_tensor(avg_loss).cuda(),
                       'train/label_loss': torch.as_tensor(avg_label_loss).cuda(),
                       'train/kd_loss': torch.as_tensor(avg_kd_loss).cuda()
                       }, sync_dist=True)

        self.training_step_outputs.clear()

    def validation_step(self, val_batch, batch_idx):
        x, _, y = val_batch
        x = self.mel_forward(x)
        y_hat, _ = self.model(x)
        loss = F.binary_cross_entropy_with_logits(y_hat, y)
        preds = torch.sigmoid(y_hat)
        results = {'val_loss': loss, "preds": preds, "targets": y}
        results = {k: v.cpu() for k, v in results.items()}
        self.validation_step_outputs.append(results)

    def on_validation_epoch_end(self):
        loss = torch.stack([x['val_loss'] for x in self.validation_step_outputs])
        preds = torch.cat([x['preds'] for x in self.validation_step_outputs], dim=0)
        targets = torch.cat([x['targets'] for x in self.validation_step_outputs], dim=0)

        all_preds = self.all_gather(preds).reshape(-1, preds.shape[-1]).cpu().float().numpy()
        all_targets = self.all_gather(targets).reshape(-1, targets.shape[-1]).cpu().float().numpy()
        all_loss = self.all_gather(loss).reshape(-1,)

        try:
            average_precision = metrics.average_precision_score(
                all_targets, all_preds, average=None)
        except ValueError:
            average_precision = np.array([np.nan] * 527)
        try:
            roc = metrics.roc_auc_score(all_targets, all_preds, average=None)
        except ValueError:
            roc = np.array([np.nan] * 527)
        logs = {'val/loss': torch.as_tensor(all_loss).mean().cuda(),
                'val/ap': torch.as_tensor(average_precision).mean().cuda(),
                'val/roc': torch.as_tensor(roc).mean().cuda()
                }
        self.log_dict(logs, sync_dist=False)
        self.validation_step_outputs.clear()


def train(config):
    # Train Models from scratch or ImageNet pre-trained on AudioSet
    # PaSST ensemble (https://github.com/kkoutini/PaSST) stored in 'resources/passt_enemble_logits_mAP_495.npy'
    # can be used as a teacher.

    # logging is done using wandb
    wandb_logger = WandbLogger(
        project="EfficientAudioTagging",
        notes="Training efficient audio tagging models on AudioSet using Knowledge Distillation.",
        tags=["AudioSet", "Audio Tagging", "Knowledge Disitillation"],
        config=config,
        name=config.experiment_name
    )

    train_dl = DataLoader(dataset=get_full_training_set(resample_rate=config.resample_rate,
                                                        roll=config.roll,
                                                        wavmix=config.wavmix,
                                                        gain_augment=config.gain_augment),
                          sampler=get_ft_weighted_sampler(config.epoch_len),  # sampler important to balance classes
                          worker_init_fn=worker_init_fn,
                          num_workers=config.num_workers,
                          batch_size=config.batch_size)

    # eval dataloader
    eval_dl = DataLoader(dataset=get_test_set(resample_rate=config.resample_rate),
                         worker_init_fn=worker_init_fn,
                         num_workers=config.num_workers,
                         batch_size=config.batch_size)

    # create pytorch lightening module
    pl_module = PLModule(config)

    # create monitor to keep track of learning rate - we want to check the behaviour of our learning rate schedule
    lr_monitor = LearningRateMonitor(logging_interval='epoch')
    # create the pytorch lightening trainer by specifying the number of epochs to train, the logger,
    # on which kind of device(s) to train and possible callbacks
    trainer = pl.Trainer(max_epochs=config.n_epochs,
                         logger=wandb_logger,
                         accelerator='auto',
                         devices=config.num_devices,
                         precision=config.precision,
                         num_sanity_val_steps=0,
                         callbacks=[lr_monitor])

    # start training and validation for the specified number of epochs
    trainer.fit(pl_module, train_dl, eval_dl)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Example of parser. ')

    # general
    parser.add_argument('--experiment_name', type=str, default="AudioSet")
    parser.add_argument('--batch_size', type=int, default=120)
    parser.add_argument('--num_workers', type=int, default=12)
    parser.add_argument('--num_devices', type=int, default=4)

    # evaluation
    # if ensemble is set, 'model_name' is not used
    parser.add_argument('--ensemble', nargs='+', default=[])
    parser.add_argument('--model_name', type=str, default="mn10_as")  # used also for training
    parser.add_argument('--cuda', action='store_true', default=False)

    # training
    parser.add_argument('--precision', type=int, default=16)
    parser.add_argument('--pretrained', action='store_true', default=False)
    parser.add_argument('--pretrain_final_temp', type=float, default=30.0)  # for DyMN
    parser.add_argument('--model_width', type=float, default=1.0)
    parser.add_argument('--strides', nargs=4, default=[2, 2, 2, 2], type=int)
    parser.add_argument('--head_type', type=str, default="mlp")
    parser.add_argument('--se_dims', type=str, default="c")
    parser.add_argument('--n_epochs', type=int, default=200)
    parser.add_argument('--mixup_alpha', type=float, default=0.3)
    parser.add_argument('--epoch_len', type=int, default=100000)
    parser.add_argument('--roll', action='store_true', default=False)
    parser.add_argument('--wavmix', action='store_true', default=False)
    parser.add_argument('--gain_augment', type=int, default=0)

    # optimizer
    parser.add_argument('--adamw', action='store_true', default=False)
    parser.add_argument('--weight_decay', type=float, default=0.0001)
    # lr schedule
    parser.add_argument('--max_lr', type=float, default=0.003)
    parser.add_argument('--warm_up_len', type=int, default=8)
    parser.add_argument('--ramp_down_start', type=int, default=80)
    parser.add_argument('--ramp_down_len', type=int, default=95)
    parser.add_argument('--last_lr_value', type=float, default=0.01)

    # knowledge distillation
    parser.add_argument('--teacher_preds', type=str,
                        default=os.path.join("resources", "passt_enemble_logits_mAP_495.npy"))
    parser.add_argument('--fname_to_index', type=str,
                        default=os.path.join("resources", "fname_to_index.pkl"))
    parser.add_argument('--temperature', type=float, default=1)
    parser.add_argument('--kd_lambda', type=float, default=0.1)

    # preprocessing
    parser.add_argument('--resample_rate', type=int, default=32000)
    parser.add_argument('--window_size', type=int, default=800)
    parser.add_argument('--hop_size', type=int, default=320)
    parser.add_argument('--n_fft', type=int, default=1024)
    parser.add_argument('--n_mels', type=int, default=128)
    parser.add_argument('--freqm', type=int, default=0)
    parser.add_argument('--timem', type=int, default=0)
    parser.add_argument('--fmin', type=int, default=0)
    parser.add_argument('--fmax', type=int, default=None)
    parser.add_argument('--fmin_aug_range', type=int, default=10)
    parser.add_argument('--fmax_aug_range', type=int, default=2000)

    args = parser.parse_args()
    train(args)