train.py

# train.py
import os, json
from omegaconf import OmegaConf
from torch.utils.tensorboard import SummaryWriter
from src.utils.seed import fix_seed
from src.utils.tb import TBLogger
from src.utils.dyn import load_obj
from src.utils.checkpoint import Checkpointer

def make_optimizer(params, conf):
    name = conf.optim.name
    print(name)
    if name == "adamw" or name == "adamW":
        import torch
        return torch.optim.AdamW(params, lr=conf.optim.lr,
                                 betas=tuple(conf.optim.betas),
                                 weight_decay=conf.optim.weight_decay)
    if name == "adam_innov":
        # Lazy import only if needed
        AdamWithInnovation = load_obj("optim.adam_innovation:AdamWithInnovation")
        return AdamWithInnovation(
            params,
            lr=conf.optim.lr, betas=tuple(conf.optim.betas),
            weight_decay=conf.optim.weight_decay,
            corr_freq=conf.optim.corr_freq, beta_y=conf.optim.beta_y,
            mode=conf.optim.innov_mode, k_init=conf.optim.k_init,
            eta_k=conf.optim.eta_k, tau=conf.optim.tau, leak=conf.optim.leak,
            kmin=conf.optim.kmin, kmax=conf.optim.kmax,
            normalize_xi=conf.optim.normalize_xi,
            exclude_bias_bn=conf.optim.exclude_bias_bn)
    if name == "sgd":
        import torch
        return torch.optim.SGD(params, lr=conf.optim.lr,
                               momentum=conf.optim.get("momentum", 0.0),
                               weight_decay=conf.optim.get("weight_decay", 0.0))
    if name == "sghmc":
        SGHMC = load_obj("src.optimizers.sghmc:SGHMC")
        return SGHMC(params, lr=conf.optim.lr, friction=conf.optim.friction,
                    temperature=conf.optim.temperature, weight_decay=conf.optim.weight_decay)
    if name == "symplectic":
        symplectic = load_obj("src.optimizers.shdg_symplectic:SHDG_SymplecticEuler")
        return symplectic(params, lr=conf.optim.lr, gamma=conf.optim.gamma, newton_refine=3, implicit_drift=True,
                      weight_decay=conf.optim.weight_decay, beta_inv=conf.optim.beta_inv, add_logdet_drift=conf.optim.add_logdet_drift)
    if name == "baoab":
        BAOAB = load_obj("src.optimizers.shdg_baoab:SHDG_BAOAB_Innovated")
        return BAOAB(params, lr=conf.optim.lr, gamma=conf.optim.gamma, beta2=conf.optim.beta2, newton_refine=3,
                      weight_decay=conf.optim.weight_decay, beta_inv=conf.optim.beta_inv, add_logdet_drift=conf.optim.add_logdet_drift)
    raise ValueError(f"Unknown optimizer {name}")

def main(cfg_path="configs/task/cifar10_resnet18.yaml", overrides=None):
    base = OmegaConf.load("configs/base.yaml")
    task_cfg = OmegaConf.load(cfg_path)
    if overrides:
        dot = OmegaConf.from_dotlist(overrides)   # <-- fix here
        task_cfg = OmegaConf.merge(task_cfg, dot)
    conf = OmegaConf.merge(base, task_cfg)

    # seed, run name, log dir
    fix_seed(conf.seed)
    run_name = (conf.get("run_name") or conf.task.name) + f"_seed{conf.seed}"
    log_dir = os.path.join(conf.log.dir, run_name)
    os.makedirs(log_dir, exist_ok=True)

    writer = SummaryWriter(log_dir)
    tblog  = TBLogger(writer, every=conf.log.every)

    # create checkpointer
    ckpt_dir = log_dir
    ckpt = Checkpointer(
        dirpath=os.path.join(ckpt_dir, "checkpoints"),
        save_last=conf.checkpoint.get("save_last", True),
        every=conf.checkpoint.get("every", 0),
        metric_key=conf.checkpoint.get("metric_key", None),
        mode=conf.checkpoint.get("mode", "max"),
        keep_n_last=conf.checkpoint.get("keep_n_last", 1),
    )

    # --- Lazy load dataloaders/model/engine from config ---
    # Datamodule: e.g., "data.cifar10:get_dataloaders"
    get_dls = load_obj(conf.task.dataloaders)
    dls = get_dls(conf.data)  # dict with train/val/test/(extras...)

    # Model builder: e.g., "models.resnet18:build_resnet18"
    build_model = load_obj(conf.task.model_builder)
    model = build_model(**(conf.task.model_kwargs or {})).to(conf.device)

    # Optimizer
    opt = make_optimizer(model.parameters(), conf)

    # NEW: build hook from YAML if enabled
    hook = None
    if getattr(conf, "innovation", None) and conf.innovation.get("enabled", False):
        HookCls = load_obj(conf.innovation.class_path)  # "hooks.innovation_damping:InnovationDamping"
        hook = HookCls(**(conf.innovation.get("kwargs", {})))

    # Resume if possible
    if args.resume:
        ckpt_data = Checkpointer.load(args.resume, model, optimizer=opt, map_location=conf.device)
        start_epoch = ckpt_data.get("epoch", 0)
        print(f"[resume] loaded {args.resume} at epoch {start_epoch}")
    # Engine: e.g., "engines.classification:run_classification"
    run_engine = load_obj(conf.task.engine)
    results = run_engine(model, opt, dls["train"], dls.get("val"), dls["test"], conf, tblog, checkpointer=ckpt, hook=hook, **{k:v for k,v in dls.items() if k not in ("train","val","test")})

    # Save config + results
    OmegaConf.save(conf, os.path.join(log_dir, "config.yaml"))
    with open(os.path.join(log_dir, "results.json"), "w") as f:
        json.dump(results, f, indent=2)
    writer.close()
    print("[OK] saved to", log_dir)


if __name__ == "__main__":
    """
    python scripts/run_many.py \
        --cfg configs/task/cifar10_c3c2fc.yaml \
        --prefix cls_cifar10 \
        --outdir out/cifar10_3c3d_sgd \
        --seeds 0 1 2 3 4 5 6 7 8 9 \
        --over task.model=3c3d optim.name=sgd optim.lr=0.01 train.epochs=100 data.batch_size=128
    
    """
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("cfg", type=str, help="Path to YAML config, e.g. configs/task/cifar10_c3c2fc.yaml")
    ap.add_argument("overrides", nargs="*", help='Overrides like key=val (e.g. "seed=0" "optim.lr=0.01")')
    ap.add_argument("--resume", type=str, default=None, help="Path to ckpt_*.pt")


    '''Examples of overrides that now work
    Scalars:
    seed=3, train.epochs=100, optim.lr=0.01, data.batch_size=128

    Bools / None:
    optim.normalize_xi=true, task.model_kwargs.dropout=null

    Lists / tuples:
    optim.betas=[0.9,0.999]

    Nested keys:
    task.model_kwargs.num_classes=10, data.aug=light

    Strings with special chars (quote them in the shell):
    task.model_builder="models.resnet18:build_resnet18"
    '''
    args = ap.parse_args()

    # call your existing main(cfg_path, overrides)
    print(args.cfg)
    print(args.overrides)
    main(cfg_path=args.cfg, overrides=args.overrides)