src/plot2eq/training/trainer.py

import math
from dataclasses import asdict
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
import wandb
from torch.amp.autocast_mode import autocast
from tqdm.auto import tqdm

from plot2eq.config import TrainConfig
from plot2eq.core.tokenizer import Tokenizer
from plot2eq.data_utils.datamodule import build_dataloaders
from plot2eq.models.core_model import Plot2EqModel
from plot2eq.training.logging_utils import create_val_predictions_table


def get_lr_scheduler(optimizer, warmup_steps, total_steps):
    def lr_lambda(current_step):
        if current_step < warmup_steps:
            return float(current_step) / float(max(1, warmup_steps))
        progress = float(current_step - warmup_steps) / float(
            max(1, total_steps - warmup_steps)
        )
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)


def train_loop(cfg: TrainConfig):

    wandb.init(project=cfg.project_name, config=asdict(cfg))

    train_loader, val_loader, vocab_size, pad_idx = build_dataloaders(
        data_dir=cfg.data_dir, batch_size=cfg.batch_size
    )

    tokenizer = Tokenizer()

    model = Plot2EqModel(
        vocab_size=vocab_size,
        pad_idx=pad_idx,
        d_model=cfg.d_model,
        nhead=cfg.nhead,
        num_enc_layers=cfg.num_enc_layers,
        num_dec_layers=cfg.num_dec_layers,
        dropout=cfg.dropout,
    ).to(cfg.device)

    model: nn.Module = torch.compile(model)  # type: ignore

    param_dict = {param_name: param for param_name, param in model.named_parameters()}

    decay = set()
    no_decay = set()
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.ndim < 2 or "bias" in name or "norm" in name or "gamma" in name:
            no_decay.add(name)
        else:
            decay.add(name)

    optim_groups = [
        {
            "params": [param_dict[pn] for pn in sorted(list(decay))],
            "weight_decay": cfg.weight_decay,
        },
        {
            "params": [param_dict[pn] for pn in sorted(list(no_decay))],
            "weight_decay": 0.0,
        },
    ]

    optimizer = optim.AdamW(optim_groups, lr=cfg.lr, betas=(0.9, 0.95))

    total_steps = len(train_loader) * cfg.epochs
    warmup_steps = int(total_steps * cfg.warmup_ratio)
    scheduler = get_lr_scheduler(optimizer, warmup_steps, total_steps)

    criterion = nn.CrossEntropyLoss(
        ignore_index=pad_idx, label_smoothing=cfg.label_smoothing
    )

    start_epoch = 1
    best_val_loss = float("inf")
    Path("checkpoints").mkdir(parents=True, exist_ok=True)

    if cfg.resume_from_checkpoint and Path(cfg.resume_from_checkpoint).exists():
        print(f"Начинаем обучение с чекпоинта: {cfg.resume_from_checkpoint}")
        checkpoint = torch.load(cfg.resume_from_checkpoint, map_location=cfg.device)

        raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
        raw_model.load_state_dict(checkpoint["model_state_dict"])  # type: ignore

        # optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        # scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

        start_epoch = checkpoint["epoch"] + 1
        best_val_loss = checkpoint["best_val_loss"]

        print(
            f"Возобновлено с эпохи {checkpoint['epoch']}. Начинаем следующую эпоху: {start_epoch}"
        )
    else:
        print("Начинаем обучение с нуля.")

    for epoch in range(start_epoch, cfg.epochs + 1):
        model.train()
        train_loss = 0.0
        total_train_samples = 0
        pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{cfg.epochs} [Train]")

        for batch_idx, (points, tokens) in enumerate(pbar):
            points, tokens = points.to(cfg.device), tokens.to(cfg.device)

            tgt_input = tokens[:, :-1]
            tgt_expected = tokens[:, 1:].contiguous().view(-1)

            optimizer.zero_grad()

            with autocast("cuda", dtype=torch.bfloat16):
                logits = model(points, tgt_input)
                logits = logits.view(-1, vocab_size)

                loss = criterion(logits, tgt_expected)

            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
            optimizer.step()
            scheduler.step()

            batch_size = points.size(0)
            train_loss += loss.item() * batch_size
            total_train_samples += batch_size

            pbar.set_postfix(
                {
                    "loss": f"{loss.item():.4f}",
                    "lr": f"{scheduler.get_last_lr()[0]:.2e}",
                }
            )

        avg_train_loss = train_loss / total_train_samples

        model.eval()
        val_loss = 0.0
        total_val_samples = 0
        correct_tokens = 0
        total_tokens = 0

        correct_sequences = 0
        total_sequences = 0
        val_table = None

        with torch.no_grad():
            pbar_val = tqdm(val_loader, desc=f"Epoch {epoch}/{cfg.epochs} [Val]")

            for batch_idx, (points, tokens) in enumerate(pbar_val):
                points, tokens = points.to(cfg.device), tokens.to(cfg.device)

                tgt_input = tokens[:, :-1]
                tgt_expected_seq = tokens[:, 1:].contiguous()

                with autocast("cuda", dtype=torch.bfloat16):
                    logits = model(points, tgt_input)
                    preds_seq = logits.argmax(dim=-1)
                    mask_seq = tgt_expected_seq != pad_idx
                    correct_sequences += (
                        ((preds_seq == tgt_expected_seq) | ~mask_seq)
                        .all(dim=-1)
                        .sum()
                        .item()
                    )
                    total_sequences += points.size(0)

                    logits_flat = logits.view(-1, vocab_size)
                    tgt_expected_flat = tgt_expected_seq.view(-1)

                    loss = criterion(logits_flat, tgt_expected_flat)

                batch_size = points.size(0)
                val_loss += loss.item() * batch_size
                total_val_samples += batch_size

                preds_flat = preds_seq.view(-1)
                mask_flat = tgt_expected_flat != pad_idx
                correct_tokens += (
                    (preds_flat[mask_flat] == tgt_expected_flat[mask_flat]).sum().item()
                )
                total_tokens += mask_flat.sum().item()

                if batch_idx == 0 and epoch % 10 == 1:
                    val_table = create_val_predictions_table(
                        model=model,
                        points=points,
                        true_tokens=tokens,
                        tokenizer=tokenizer,
                    )

        avg_val_loss = val_loss / total_val_samples
        val_tok_acc = correct_tokens / total_tokens
        val_seq_acc = correct_sequences / total_sequences

        print(
            f"\nEpoch {epoch} | Train Loss: {avg_train_loss:.4f} | "
            f"Val Loss: {avg_val_loss:.4f} | "
            f"Token Acc: {val_tok_acc:.4f} | "
            f"Seq Acc (Exact Match): {val_seq_acc:.4f}\n"
        )

        wandb.log(
            {
                "train_loss": avg_train_loss,
                "val_loss": avg_val_loss,
                "val_tok_acc": val_tok_acc,
                "val_seq_acc": val_seq_acc,
                "learning_rate": scheduler.get_last_lr()[0],
                "epoch": epoch,
                **({"Predictions": val_table} if val_table is not None else {}),
            }
        )

        raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
        checkpoint = {
            "epoch": epoch,
            "model_state_dict": raw_model.state_dict(),  # type: ignore
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "best_val_loss": best_val_loss,
        }
        torch.save(checkpoint, "checkpoints/last_model.pth")

        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            checkpoint["best_val_loss"] = best_val_loss
            torch.save(checkpoint, "checkpoints/best_model.pth")