train_pupum2d.py

import os
import json5
import numpy as np
import random
import torch
import torch.nn.functional as F
import argparse
import math
from tqdm import tqdm
from dataloader_webdataset import StreamingAudioWebDataset

from model import PupuM2D
from util import load_config
from librosa.filters import mel as librosa_mel_fn

import accelerate
from torch.optim import AdamW
from accelerate import DistributedDataParallelKwargs
from accelerate.utils import ProjectConfiguration
from accelerate.logging import get_logger


class WarmupCosineSchedule(object):
    def __init__(
        self,
        optimizer,
        warmup_steps,
        start_lr,
        ref_lr,
        T_max,
        final_lr=0.0,
    ):
        self.optimizer = optimizer
        self.start_lr = start_lr
        self.ref_lr = ref_lr
        self.final_lr = final_lr
        self.warmup_steps = warmup_steps
        self.T_max = T_max - warmup_steps
        self._step = 0.0

    def step(self):
        self._step += 1
        if self._step < self.warmup_steps:
            progress = float(self._step) / float(max(1, self.warmup_steps))
            new_lr = self.start_lr + progress * (self.ref_lr - self.start_lr)
        else:
            progress = float(self._step - self.warmup_steps) / float(max(1, self.T_max))
            new_lr = max(
                self.final_lr,
                self.final_lr
                + (self.ref_lr - self.final_lr)
                * 0.5
                * (1.0 + math.cos(math.pi * progress)),
            )

        for group in self.optimizer.param_groups:
            group["lr"] = new_lr

        return new_lr


class CosineWDSchedule(object):
    def __init__(self, optimizer, ref_wd, T_max, final_wd=0.0):
        self.optimizer = optimizer
        self.ref_wd = ref_wd
        self.final_wd = final_wd
        self.T_max = T_max
        self._step = 0.0

    def step(self):
        self._step += 1
        progress = self._step / self.T_max
        new_wd = self.final_wd + (self.ref_wd - self.final_wd) * 0.5 * (
            1.0 + math.cos(math.pi * progress)
        )

        if self.final_wd <= self.ref_wd:
            new_wd = max(self.final_wd, new_wd)
        else:
            new_wd = min(self.final_wd, new_wd)

        for group in self.optimizer.param_groups:
            if ("WD_exclude" not in group) or not group["WD_exclude"]:
                group["weight_decay"] = new_wd
        return new_wd


class EMASchedule:
    def __init__(self, start_ema, end_ema, total_steps):
        self.start_ema = start_ema
        self.end_ema = end_ema
        self.total_steps = total_steps
        self.curr_step = 0

    def step(self):
        if self.curr_step < self.total_steps:
            m = self.start_ema + (self.end_ema - self.start_ema) * (
                self.curr_step / self.total_steps
            )
        else:
            m = self.end_ema
        self.curr_step += 1
        return m


def get_pupum2d_param_groups(model):
    regular_params = []
    no_decay_params = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue

        if (
            param.ndim <= 1
            or "bias" in name
            or "norm" in name
            or "token" in name
            or "embedding" in name
            or "pos_embed" in name
        ):
            no_decay_params.append(param)
        else:
            regular_params.append(param)

    return [
        {"params": regular_params, "WD_exclude": False},
        {"params": no_decay_params, "weight_decay": 0.0, "WD_exclude": True},
    ]


def pupum2d_loss_fn(pred, target, norm_target: bool = True):
    if norm_target:
        mean = target.mean(dim=-1, keepdim=True)
        var = target.var(dim=-1, keepdim=True)
        target = (target - mean) / (var + 1e-6).sqrt()
    loss = F.smooth_l1_loss(pred, target)
    return loss


@torch.no_grad()
def update_ema_teacher(student_params, teacher_params, momentum):
    for param_s, param_t in zip(student_params, teacher_params):
        param_t.data.mul_(momentum).add_(param_s.data, alpha=1.0 - momentum)


def force_zero_grad(optimizer):
    optimizer.zero_grad()
    inner_optimizer = getattr(optimizer, "optimizer", None)
    if inner_optimizer is not None:
        try:
            inner_optimizer.zero_grad(set_to_none=True)
        except TypeError:
            inner_optimizer.zero_grad()


mel_basis = {}
hann_window = {}


@torch.no_grad()
def extract_mel_features(y, cfg, center=False):
    global mel_basis, hann_window
    if cfg.preprocess.fmax not in mel_basis:
        mel = librosa_mel_fn(
            sr=cfg.preprocess.sample_rate,
            n_fft=cfg.preprocess.n_fft,
            n_mels=cfg.preprocess.n_mels,
            fmin=cfg.preprocess.fmin,
            fmax=cfg.preprocess.fmax,
        )
        mel_basis[str(cfg.preprocess.fmax) + "_" + str(y.device)] = (
            torch.from_numpy(mel).float().to(y.device)
        )
        hann_window[str(y.device)] = torch.hann_window(cfg.preprocess.win_size).to(
            y.device
        )

    y = torch.nn.functional.pad(
        y.unsqueeze(1),
        (
            int((cfg.preprocess.n_fft - cfg.preprocess.hop_size) / 2),
            int((cfg.preprocess.n_fft - cfg.preprocess.hop_size) / 2),
        ),
        mode="reflect",
    )
    y = y.squeeze(1)

    spec = torch.stft(
        y,
        cfg.preprocess.n_fft,
        hop_length=cfg.preprocess.hop_size,
        win_length=cfg.preprocess.win_size,
        window=hann_window[str(y.device)],
        center=center,
        pad_mode="reflect",
        normalized=False,
        onesided=True,
        return_complex=True,
    )
    spec = torch.view_as_real(spec)
    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))

    spec = torch.matmul(mel_basis[str(cfg.preprocess.fmax) + "_" + str(y.device)], spec)

    spec = torch.log(torch.clamp(spec, min=1e-5))

    if cfg.preprocess.normalize:
        mean = -4.089994845986366
        std = 2.0242277159094813
        spec = (spec - mean) / (std + 1e-8)

    if cfg.preprocess.flip_ft:
        spec = spec.transpose(-2, -1)

    return spec.unsqueeze(1)


def cuda_relevant():
    torch.cuda.empty_cache()
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True


def dump_cfg(cfg, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    json5.dump(
        cfg,
        open(path, "w"),
        indent=4,
        sort_keys=True,
        ensure_ascii=False,
        quote_keys=True,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config.json", required=True)
    parser.add_argument("--exp_name", type=str, required=True)
    parser.add_argument("--resume_type", type=str)
    parser.add_argument("--checkpoint", type=str)
    parser.add_argument("--log_level", default="warning")
    args = parser.parse_args()
    cfg = load_config(args.config)
    cuda_relevant()

    cfg.exp_name = args.exp_name

    ddp_kwargs = DistributedDataParallelKwargs()
    exp_dir = os.path.join(os.path.abspath(cfg.log_dir), args.exp_name)
    project_config = ProjectConfiguration(
        project_dir=exp_dir, logging_dir=os.path.join(exp_dir, "log")
    )
    accelerator = accelerate.Accelerator(
        gradient_accumulation_steps=cfg.train.gradient_accumulation_step,
        log_with=cfg.train.tracker,
        project_config=project_config,
        kwargs_handlers=[ddp_kwargs],
    )

    if accelerator.is_main_process:
        os.makedirs(project_config.project_dir, exist_ok=True)
        os.makedirs(project_config.logging_dir, exist_ok=True)

    with accelerator.main_process_first():
        accelerator.init_trackers(args.exp_name)
        logger = get_logger(args.exp_name, log_level=args.log_level)

    checkpoint_dir = os.path.join(exp_dir, "checkpoint")
    if accelerator.is_main_process:
        os.makedirs(checkpoint_dir, exist_ok=True)

    step = 0

    with accelerator.main_process_first():
        random.seed(cfg.train.random_seed)
        np.random.seed(cfg.train.random_seed)
        torch.random.manual_seed(cfg.train.random_seed)

    with accelerator.main_process_first():
        train_dataset = StreamingAudioWebDataset(
            "TBD", cfg
        )
        train_dataloader = train_dataset.get_dataloader()

    with accelerator.main_process_first():
        model = PupuM2D(cfg)

    with accelerator.main_process_first():
        param_groups = get_pupum2d_param_groups(model)
        optimizer = AdamW(
            param_groups,
            lr=cfg.train.optimizer.start_lr,
            betas=(cfg.train.optimizer.adam_b1, cfg.train.optimizer.adam_b2),
            eps=1e-6,
        )
        lr_scheduler = WarmupCosineSchedule(
            optimizer,
            warmup_steps=cfg.train.warmup_steps,
            start_lr=cfg.train.optimizer.start_lr,
            ref_lr=cfg.train.optimizer.peak_lr,
            final_lr=cfg.train.optimizer.final_lr,
            T_max=cfg.train.total_steps,
        )
        wd_scheduler = CosineWDSchedule(
            optimizer,
            ref_wd=cfg.train.optimizer.start_wd,
            final_wd=cfg.train.optimizer.final_wd,
            T_max=cfg.train.total_steps,
        )

    ema_scheduler = EMASchedule(
        start_ema=cfg.train.optimizer.start_ema,
        end_ema=cfg.train.optimizer.final_ema,
        total_steps=cfg.train.total_steps,
    )

    train_dataloader, model, optimizer = accelerator.prepare(
        train_dataloader, model, optimizer
    )

    with accelerator.main_process_first():
        if args.resume_type == "resume":
            accelerator.load_state(args.checkpoint)
            try:
                step = int(args.checkpoint.split("step-")[-1].split("_")[0]) + 1
                if step > 0:
                    logger.info(
                        f"Resumed from step {step}, Fast-forwarding schedulers..."
                    )
                    lr_scheduler._step = float(step - 1)
                    wd_scheduler._step = float(step - 1)
                    ema_scheduler.curr_step = step - 1
                    lr_scheduler.step()
                    wd_scheduler.step()
            except:
                logger.warning("Could not parse epoch/step from checkpoint path")

        config_save_path = os.path.join(exp_dir, "args.json")
        if accelerator.is_main_process:
            dump_cfg(cfg, config_save_path)

    model.train()
    optimizer.zero_grad()
    train_iter = iter(train_dataloader)

    progress_bar = tqdm(
        range(step, cfg.train.total_steps),
        desc="Training",
        disable=not accelerator.is_main_process,
        initial=step,
        total=cfg.train.total_steps,
        dynamic_ncols=True,
    )

    curr_lr = optimizer.param_groups[0]["lr"]
    curr_wd = cfg.train.optimizer.start_wd
    momentum = cfg.train.optimizer.start_ema
    accumulated_loss = 0.0

    while step < cfg.train.total_steps:
        try:
            batch = next(train_iter)
        except StopIteration:
            train_iter = iter(train_dataloader)
            batch = next(train_iter)
        except Exception as e:
            logger.warning(f"Batch read error: {e}")
            continue

        with accelerator.accumulate(model):
            audio_gt = batch.to(accelerator.device).squeeze()
            mel_gt = extract_mel_features(audio_gt, cfg)

            with accelerator.autocast():
                z_pred, h_target = model(mel_gt, step)
                loss = pupum2d_loss_fn(
                    z_pred, h_target, norm_target=cfg.train.norm_pix_loss
                )

            if torch.isnan(loss) or torch.isinf(loss):
                logger.error(f"Loss NaN/Inf at step {step}!")
                force_zero_grad(optimizer)
                accumulated_loss = 0.0
                continue

            accumulated_loss += loss.item() / cfg.train.gradient_accumulation_step
            accelerator.backward(loss)

            if accelerator.sync_gradients:
                accelerator.clip_grad_norm_(model.parameters(), cfg.train.clip_grad)

            optimizer.step()
            optimizer.zero_grad()

            if not accelerator.sync_gradients:
                continue

            curr_lr = lr_scheduler.step()
            curr_wd = wd_scheduler.step()

            unwrapped_model = accelerator.unwrap_model(model)
            momentum = ema_scheduler.step()
            update_ema_teacher(
                unwrapped_model.student.parameters(),
                unwrapped_model.teacher.parameters(),
                momentum,
            )

            loss_val = accumulated_loss
            accumulated_loss = 0.0

            if step % 100 == 0:
                accelerator.log(
                    {
                        "Train/Loss": loss_val,
                        "Train/LR": curr_lr,
                        "Train/WD": curr_wd,
                        "Train/EMA": momentum,
                    },
                    step=step,
                )
                if accelerator.is_main_process:
                    progress_bar.set_postfix(
                        {
                            "loss": f"{loss_val:.4f}",
                            "lr": f"{curr_lr:.2e}",
                            "wd": f"{curr_wd:.3f}",
                        }
                    )

            step = step + 1
            progress_bar.update(1)

            if step % cfg.train.save_checkpoint_stride == 0 and step > 0:
                if accelerator.is_main_process:
                    save_path = os.path.join(
                        checkpoint_dir,
                        "step-{:07d}_loss-{:.6f}".format(step, loss_val),
                    )
                    accelerator.save_state(save_path)

    accelerator.end_training()