main.py

# KAN quantization
# Initial code was adopted from: github.com/zhutmost/lsq-net

import logging
from pathlib import Path
import math
import torch
import torch.nn as nn
import yaml
import pdb
import process
import qat
import util
from model import create_model
from qat.quantizer.quantizer import Quantizer
from util.checkpoint import EarlyStopping, _load_checkpoint_compat, _save_checkpoint_compat
from util.optimizers import _get, build_optimizer
from util.lr_scheduler import build_lr_scheduler_from_cfg
from model.vgg_kan_imagenet import *


def log_quantizer_stats(model, writer, global_step: int):
    """
    Logs quantizer stats safely:
      - numbers / 0-d tensors -> add_scalar
      - strings -> add_text
      - everything else -> best-effort convert or skip
    """
    
    # unwrap DataParallel if present
    if isinstance(model, nn.DataParallel):
        model = model.module

    def _to_float(v):
        # returns float or None if cannot be represented as scalar
        if v is None:
            return None

        if isinstance(v, bool):
            return float(v)

        if isinstance(v, (int, float)):
            if not math.isfinite(float(v)):
                return None
            return float(v)

        if isinstance(v, torch.Tensor):
            if v.numel() != 1:
                return None
            x = v.detach()
            # move to cpu safely
            x = x.float().cpu().item()
            if not math.isfinite(float(x)):
                return None
            return float(x)

        # try last-resort conversion (e.g., numpy scalar)
        try:
            x = float(v)
            if not math.isfinite(float(x)):
                return None
            return float(x)
        except Exception:
            return None

    for name, module in model.named_modules():
        if isinstance(module, Quantizer) and getattr(module, "record_stats", False):
            stats = module.export_stats()
            if not stats:
                continue

            for k, v in stats.items():
                tag = f"QStats/{name}/{k}"

                if isinstance(v, str):
                    # strings cannot be logged as scalars
                    writer.add_text(tag, v, global_step)
                    continue

                fv = _to_float(v)
                if fv is None:
                    # skip non-scalars / non-finite values
                    continue
                writer.add_scalar(tag, fv, global_step)


def main():
    script_dir = Path.cwd()
    args = util.get_config(default_file=script_dir / "config.yaml")
    output_dir = script_dir / args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    # Eval requires a checkpoint path
    if args.eval:
        assert args.resume.path != "", "For evaluation, resume.path must point to a trained model checkpoint."

    # Logger
    log_dir = util.init_logger(args.name, output_dir, script_dir / "logging.conf")
    logger = logging.getLogger()

    # dump experiment config (munch -> dict)
    try:
        cfg_dump = args.toDict()
    except Exception:
        cfg_dump = dict(args)
    with open(log_dir / "args.yaml", "w") as yaml_file:
        yaml.safe_dump(cfg_dump, yaml_file)

    pymonitor = util.ProgressMonitor(logger)
    tbmonitor = util.TensorBoardMonitor(logger, log_dir)
    monitors = [pymonitor, tbmonitor]

    # Device sanity
    if args.device.type == "cpu" or (not torch.cuda.is_available()) or args.device.gpu == []:
        args.device.gpu = []
    else:
        available_gpu = torch.cuda.device_count()
        for dev_id in args.device.gpu:
            if dev_id >= available_gpu:
                logger.error("GPU device ID %d requested, but only %d devices available", dev_id, available_gpu)
                raise SystemExit(1)
        torch.cuda.set_device(args.device.gpu[0])
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False

    # Data
    train_loader, val_loader, test_loader = util.load_data(args.dataloader)
    logger.info(
        "Dataset `%s` size:\n"
        "          Training Set = %d (%d)\n"
        "        Validation Set = %d (%d)\n"
        "              Test Set = %d (%d)",
        args.dataloader.dataset,
        len(train_loader.sampler), len(train_loader),
        len(val_loader.sampler), len(val_loader),
        len(test_loader.sampler), len(test_loader),
    )

    # ----------------------------
    # Build model (MUST happen before resume)
    # ----------------------------
    model = create_model(args)
    print(model)

    # Load FP32 pretrained checkpoint into the plain model (weights only)
    if getattr(args.pretrained, "pretrain_from_fp32", False):
        model, _, _ = _load_checkpoint_compat(
            model,
            args.pretrained.pretrained_fp32_path,
            args.device.type,
            lean=True,
            strict=True,
            restore_rng_state=False,
        )

    # HF ImageNet weights (only if imagenet + supported arch)
    if args.dataloader.dataset == "imagenet" and getattr(args.pretrained, "load_from_hf", False):
        if args.arch == "kagn_v2":
            model, ckpt_name, missing, unexpected, _ = util.load_hf_weights_into_model(
                model, repo_id="brivangl/vgg_kagn11_v2", device="cpu"
            )
        elif args.arch == "kagn_v4":
            model, ckpt_name, missing, unexpected, _ = util.load_hf_weights_into_model(
                model, repo_id="brivangl/vgg_kagn11_v4", device="cpu"
            )
    # HF ImageNet weights (only if imagenet + supported arch)
    if args.dataloader.dataset == "tinyimagenet" and getattr(args.pretrained, "load_from_hf", False):
        if args.arch == "kagn_v2":
            model = imagenet_kagn_v4_loaded_for_tinyimagenet()
        elif args.arch == "kagn_v4":
            model = imagenet_kagn_v4_loaded_for_tinyimagenet()
        
        logger.info("The model is overloaded with imagenet model weights!")

        # pdb.set_trace()


    # Optional: convert NN -> KAN
    if getattr(args.kan, "kan_convert", False):
        modules_to_replace = qat.find_modules_to_kan_from_nn(
            model, args.quan,
            kan=args.kan.kan_convert,
            kan_variant=args.kan.kan_variant
        )
        model = qat.replace_module_by_names(model, modules_to_replace)

    # Optional: insert quantizers
    if getattr(args.quan, "quantization", False):
        modules_to_quantize = qat.find_modules_to_quantize(model, args.quan)
        model = qat.replace_module_by_names(model, modules_to_quantize)
        logger.info("Inserted quantizers into the original model")

    print(model)
    qat.print_quant_coverage(model)
    util.count_parameters(model)

    logger.info("The dataset will be trained on this model:")
    logger.info(model)

    # DataParallel
    if args.device.gpu and not getattr(args.dataloader, "serialized", False):
        model = torch.nn.DataParallel(model, device_ids=args.device.gpu)

    model.to(args.device.type)

    # Loss
    criterion = torch.nn.CrossEntropyLoss().to(args.device.type)

    # ----------------------------
    # Optimizer + Scheduler + AMP scaler (MUST exist before true resume)
    # ----------------------------
    ###################################################################################################
    optimizer = build_optimizer(model, args.optimizer)

    lr_scheduler = build_lr_scheduler_from_cfg(
    optimizer,
    args.lr_scheduler,
    batch_size=train_loader.batch_size,
    num_samples=len(train_loader.sampler),
    accum_steps=int(getattr(args.optimizer, "accum_steps", 1)),epochs=(int(getattr(args, "epochs")) if getattr(args, "epochs", None) is not None else None),)

    global_step = 0
    # scaler = torch.amp.GradScaler("cuda") if args.device.type == "cuda" else None
    amp_cfg = getattr(args, "amp", None)
    use_amp = bool(_get(amp_cfg, "enable", False)) and (args.device.type == "cuda")
    use_grad_scaler = bool(_get(amp_cfg, "grad_scaler", True)) and use_amp

    scaler = torch.amp.GradScaler("cuda") if use_grad_scaler else None
    ###################################################################################################

    logger.info(("Optimizer: %s" % optimizer).replace("\n", "\n" + " " * 11))
    logger.info("LR scheduler: %s\n", lr_scheduler)

    perf_scoreboard = process.PerformanceScoreboard(args.log.num_best_scores)

    # Early stopping
    es_cfg = getattr(args, "early_stopping", None)
    early_stopper = None
    if es_cfg is not None and getattr(es_cfg, "enable", False):
        early_stopper = EarlyStopping(
            mode=getattr(es_cfg, "mode", "max"),
            patience=getattr(es_cfg, "patience", 20),
            min_delta=getattr(es_cfg, "min_delta", 0.0),
            warmup_epochs=getattr(es_cfg, "warmup_epochs", 0),
        )
        logger.info(
            "EarlyStopping enabled: monitor=%s, mode=%s, patience=%s, min_delta=%s, warmup_epochs=%s",
            getattr(es_cfg, "monitor", "val_top1"),
            getattr(es_cfg, "mode", "max"),
            getattr(es_cfg, "patience", 20),
            getattr(es_cfg, "min_delta", 0.0),
            getattr(es_cfg, "warmup_epochs", 0),
        )

    # ----------------------------
    # TRUE RESUME (only once, after everything exists)
    # ----------------------------
    start_epoch = 0
    resume_extras = None
    if args.resume.path:
        model, start_epoch, resume_extras = _load_checkpoint_compat(
            model,
            args.resume.path,
            args.device.type,
            lean=args.resume.lean,
            optimizer=None if args.resume.lean else optimizer,
            lr_scheduler=None if args.resume.lean else lr_scheduler,
            scaler=None if args.resume.lean else scaler,
            strict=False,
            restore_rng_state=(not args.eval and not args.resume.lean),
        )

        if isinstance(resume_extras, dict):
            if resume_extras.get("global_step") is not None:
                try:
                    global_step = int(resume_extras.get("global_step"))
                except Exception:
                    pass

            # scoreboard restore
            if resume_extras.get("scoreboard") is not None:
                try:
                    perf_scoreboard.board = list(resume_extras["scoreboard"])
                except Exception:
                    pass

            # early stopper restore
            if early_stopper is not None and resume_extras.get("early_stopper") is not None:
                try:
                    early_stopper.load_state_dict(resume_extras["early_stopper"])
                except Exception:
                    pass

        # fallback global_step if old ckpt
        if (not resume_extras) or ("global_step" not in resume_extras):
            try:
                global_step = int(start_epoch) * int(len(train_loader))
            except Exception:
                pass

    # ----------------------------
    # Eval or Train
    # ----------------------------
    if args.eval:
        process.validate(test_loader, model, criterion, -1, monitors, args)
        tbmonitor.writer.close()
        logger.info("Program completed successfully ... exiting ...")
        logger.info("===============================================")
        return

    # optional: pre-trained evaluation (safe getattr)
    pre_trained_flag = bool(getattr(args, "pre_trained", False))
    if args.resume.path or pre_trained_flag:
        logger.info(">>>>>>>> Epoch -1 (pre-trained model evaluation)")
        top1, top5, _ = process.validate(val_loader, model, criterion, start_epoch - 1, monitors, args)
        perf_scoreboard.update(top1, top5, start_epoch - 1)

    for epoch in range(start_epoch, args.epochs):
        logger.info(">>>>>>>> Epoch %3d", epoch)

        t_top1, t_top5, t_loss, global_step = process.train(
            train_loader, model, criterion, optimizer,
            lr_scheduler, epoch, monitors, args,
            scaler=scaler, global_step=global_step
        )

        v_top1, v_top5, v_loss = process.validate(val_loader, model, criterion, epoch, monitors, args)

        tbmonitor.writer.add_scalars("Train_vs_Validation/Loss", {"train": t_loss, "val": v_loss}, epoch)
        tbmonitor.writer.add_scalars("Train_vs_Validation/Top1", {"train": t_top1, "val": v_top1}, epoch)
        tbmonitor.writer.add_scalars("Train_vs_Validation/Top5", {"train": t_top5, "val": v_top5}, epoch)

        perf_scoreboard.update(v_top1, v_top5, epoch)
        is_best = perf_scoreboard.is_best(epoch)

        # Save checkpoint (supports both old + new signatures)
        _save_checkpoint_compat(
            epoch,
            args.arch,
            model,
            {"top1": v_top1, "top5": v_top5},
            is_best,
            args.name,
            log_dir,
            optimizer=optimizer,
            lr_scheduler=lr_scheduler,
            scaler=scaler,
            global_step=global_step,
            early_stopper=early_stopper,
            scoreboard=perf_scoreboard.board,
        )

        # Log quantizer stats (use global_step so it aligns with resume)
        log_quantizer_stats(model, tbmonitor.writer, global_step)

        # Early stopping
        if early_stopper is not None:
            monitor = getattr(args.early_stopping, "monitor", "val_top1")
            metric = v_loss if monitor == "val_loss" else v_top1

            should_stop = early_stopper.step(metric, epoch)
            logger.info(
                "[EarlyStop] monitor=%s metric=%.6f best=%s best_epoch=%s bad=%d/%d",
                monitor, float(metric),
                str(early_stopper.best), str(early_stopper.best_epoch),
                early_stopper.num_bad, early_stopper.patience
            )
            if should_stop:
                logger.info(
                    ">>>>>>> Early stopping triggered at epoch %d. Best epoch was %s.",
                    epoch, str(early_stopper.best_epoch)
                )
                break

    # Final: evaluate best checkpoint if available
    logger.info(">>>>>>>> Final Model evaluation (BEST checkpoint): ")

    best_ckpt = log_dir / f"{args.name}_best.pth.tar"
    best_epoch = None

    if best_ckpt.exists():
        ckpt = torch.load(str(best_ckpt), map_location="cpu")
        best_epoch = ckpt.get("epoch", None)

        model, _, _ = _load_checkpoint_compat(
            model,
            str(best_ckpt),
            args.device.type,
            lean=True,
            strict=False,
            restore_rng_state=False,
        )

        logger.info("Loaded BEST checkpoint: %s (epoch=%s)", str(best_ckpt), str(best_epoch))
    else:
        logger.warning("BEST checkpoint not found: %s. Falling back to last epoch weights.", str(best_ckpt))

    process.validate(test_loader, model, criterion, best_epoch if best_epoch is not None else -1, monitors, args)

    tbmonitor.writer.close()
    logger.info("Program completed successfully ... exiting ...")
    logger.info("===============================================")


if __name__ == "__main__":
    main()



# #### changes made:
# # 1. main.py -> # tbmonitor.writer.add_graph(model, input_to_model=train_loader.dataset[0][0].unsqueeze(0)) #  not working
# # 2. quan/func needs bias, changed
# # 3. process.py needs a new accuracy function 
# # 4. cifar10_model.py needs l-124 out = F.avg_pool2d(out, (out.size()[3], out.size()[3])) ### modified   



# tensorboard --logdir out/KAN_MLP_MNIST_LSQ_W4A4/tb_runs