scripts/train_baseline.py

from __future__ import annotations

import argparse
import json
import random
import time
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import torch
from tqdm import tqdm

import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from data.io_utils import write_json  # noqa: E402
from train.config_utils import (  # noqa: E402
    apply_dot_overrides,
    deep_merge,
    load_yaml,
    parse_overrides_json,
)
from train.dataset import (  # noqa: E402
    ChairAttributeDataset,
    build_dataloader,
    class_pos_weights,
    load_json,
    split_paths,
)
from train.encoders import build_backbones  # noqa: E402
from train.losses import bce_logits_loss  # noqa: E402
from train.metrics import compute_metrics  # noqa: E402
from train.models import BaselineClassifier  # noqa: E402


def log(message: str) -> None:
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{ts}] [train_baseline] {message}", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train baseline model (A/B/C).")
    parser.add_argument("--dataset-config", type=Path, default=REPO_ROOT / "configs" / "dataset.yaml")
    parser.add_argument("--model-config", type=Path, default=REPO_ROOT / "configs" / "model.yaml")
    parser.add_argument("--train-config", type=Path, default=REPO_ROOT / "configs" / "train.yaml")
    parser.add_argument("--eval-config", type=Path, default=REPO_ROOT / "configs" / "eval.yaml")
    parser.add_argument("--run-name", type=str, required=True)
    parser.add_argument("--overrides-json", type=str, default="")
    parser.add_argument("--mode", type=str, default="full", choices=["smoke", "full"])
    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"])
    return parser.parse_args()


def set_seed(seed: int, deterministic: bool) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = deterministic
    torch.backends.cudnn.benchmark = not deterministic


def load_combined_config(args: argparse.Namespace) -> Dict[str, Any]:
    dataset_cfg = load_yaml(args.dataset_config)
    model_cfg = load_yaml(args.model_config)
    train_cfg = load_yaml(args.train_config)
    eval_cfg = load_yaml(args.eval_config)
    combined = deep_merge(dataset_cfg, model_cfg)
    combined = deep_merge(combined, train_cfg)
    combined = deep_merge(combined, eval_cfg)
    overrides = parse_overrides_json(args.overrides_json)
    if overrides:
        combined = apply_dot_overrides(combined, overrides)
    return combined


def build_loaders(cfg: Dict[str, Any], batch_size: int, eval_batch_size: int, smoke: bool) -> Tuple[Any, Any, Any, Dict[str, int]]:
    processed_root = REPO_ROOT / cfg["dataset"]["processed_dir"]
    label_vocab_path = processed_root / cfg["processing"]["label_vocab_file"]
    label_freq_path = processed_root / cfg["processing"]["label_frequency_file"]
    cache_dir = REPO_ROOT / cfg["dataset"]["root_dir"] / "images"
    splits = split_paths(processed_root=processed_root, split_dir_name=cfg["processing"]["split_dir"])

    label_vocab = load_json(label_vocab_path)
    label_freq = load_json(label_freq_path)

    train_ds = ChairAttributeDataset(splits["train"], label_vocab_path, REPO_ROOT, cache_dir)
    val_ds = ChairAttributeDataset(splits["val"], label_vocab_path, REPO_ROOT, cache_dir)
    test_ds = ChairAttributeDataset(splits["test"], label_vocab_path, REPO_ROOT, cache_dir)

    num_workers = int(cfg["dataset"].get("num_workers", 0))
    if smoke:
        num_workers = min(num_workers, 8)

    train_loader = build_dataloader(
        dataset=train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
    )
    val_loader = build_dataloader(
        dataset=val_ds,
        batch_size=eval_batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
    )
    test_loader = build_dataloader(
        dataset=test_ds,
        batch_size=eval_batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
    )
    return train_loader, val_loader, test_loader, {"vocab": label_vocab, "freq": label_freq}


def build_model(cfg: Dict[str, Any], num_labels: int, device: torch.device) -> BaselineClassifier:
    use_clip = bool(cfg["encoders"]["clip"]["enabled"])
    use_dino = bool(cfg["encoders"]["dino"]["enabled"])
    clip_name = str(cfg["encoders"]["clip"].get("model_name", "vit_base_patch16_clip_224.openai"))
    dino_name = str(cfg["encoders"]["dino"].get("model_name", "vit_base_patch14_dinov2.lvd142m"))
    if clip_name == "ViT-B-32":
        clip_name = "vit_base_patch16_clip_224.openai"
    if dino_name == "vit_base_patch14_dinov2":
        dino_name = "vit_base_patch14_dinov2.lvd142m"

    clip_backbone, dino_backbone, feature_dim = build_backbones(
        use_clip=use_clip,
        use_dino=use_dino,
        clip_model_name=clip_name,
        dino_model_name=dino_name,
        clip_pretrained=bool(cfg["encoders"]["clip"].get("pretrained", True)),
        dino_pretrained=bool(cfg["encoders"]["dino"].get("pretrained", True)),
    )
    hidden_dims = list(cfg["baseline_heads"]["mlp"].get("hidden_dims", [512, 256]))
    dropout = float(cfg["baseline_heads"]["mlp"].get("dropout", 0.2))
    model = BaselineClassifier(
        clip_backbone=clip_backbone,
        dino_backbone=dino_backbone,
        feature_dim=feature_dim,
        hidden_dims=hidden_dims,
        dropout=dropout,
        num_labels=num_labels,
    )
    model.freeze_backbones()
    return model.to(device)


def run_eval(model: BaselineClassifier, loader: Any, device: torch.device, pos_weight: torch.Tensor) -> Dict[str, Any]:
    model.eval()
    all_logits: List[np.ndarray] = []
    all_targets: List[np.ndarray] = []
    losses: List[float] = []
    with torch.no_grad():
        for batch in loader:
            targets = batch["targets"].to(device)
            logits = model(batch["images"], device=device)
            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
            losses.append(float(loss.item()))
            all_logits.append(logits.detach().cpu().numpy())
            all_targets.append(targets.detach().cpu().numpy())
    logits_np = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, 1))
    targets_np = np.concatenate(all_targets, axis=0) if all_targets else np.zeros((0, 1))
    metrics = compute_metrics(targets=targets_np, logits=logits_np)
    metrics["loss"] = float(np.mean(losses)) if losses else 0.0
    return metrics


def train() -> None:
    args = parse_args()
    log(f"Starting run_name={args.run_name} mode={args.mode} device_request={args.device}")
    log(f"dataset_config={args.dataset_config}")
    log(f"model_config={args.model_config}")
    log(f"train_config={args.train_config}")
    cfg = load_combined_config(args)
    run_cfg = cfg.get("run", {})
    seed = int(run_cfg.get("seed", 42))
    deterministic = bool(run_cfg.get("deterministic", True))
    set_seed(seed=seed, deterministic=deterministic)
    log(f"Seed set to {seed}, deterministic={deterministic}")

    smoke = args.mode == "smoke"
    train_epochs = 2 if smoke else int(cfg["training"]["stage1"].get("epochs", 20))
    batch_size = 4 if smoke else int(cfg["batching"].get("batch_size", 32))
    eval_batch_size = 4 if smoke else int(cfg["batching"].get("eval_batch_size", 64))

    cuda_available = torch.cuda.is_available()
    selected_device = "cuda" if (args.device == "auto" and cuda_available) else args.device
    if args.device == "auto" and not cuda_available:
        selected_device = "cpu"
    if selected_device == "cuda" and not cuda_available:
        raise RuntimeError(
            "CUDA requested but not available in torch build/environment. "
            "Install CUDA-enabled PyTorch in the active virtualenv."
        )
    device = torch.device(selected_device)
    log(
        f"Resolved device={device.type}, torch_cuda_available={cuda_available}, "
        f"torch_version={torch.__version__}, torch_cuda={torch.version.cuda}"
    )
    log("Building dataloaders from processed manifests...")
    train_loader, val_loader, test_loader, label_info = build_loaders(
        cfg=cfg,
        batch_size=batch_size,
        eval_batch_size=eval_batch_size,
        smoke=smoke,
    )
    log(
        f"Dataloader sizes -> train_batches={len(train_loader)} val_batches={len(val_loader)} "
        f"test_batches={len(test_loader)}"
    )
    label_vocab = label_info["vocab"]
    label_freq = label_info["freq"]
    num_labels = len(label_vocab)
    log(f"Label space size: {num_labels}")
    log("Building encoder backbones (first run may download pretrained weights)...")
    model = build_model(cfg=cfg, num_labels=num_labels, device=device)
    log("Model construction complete")
    pos_weight = class_pos_weights(label_frequencies=label_freq, label_vocab=label_vocab).to(device)

    lr = float(cfg["training"]["stage1"].get("lr_head", 1e-3))
    weight_decay = float(cfg["optimization"].get("weight_decay", 1e-4))
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=lr,
        weight_decay=weight_decay,
    )

    output_root = REPO_ROOT / run_cfg.get("output_dir", "outputs")
    run_dir = output_root / args.run_name
    run_dir.mkdir(parents=True, exist_ok=True)

    history: List[Dict[str, Any]] = []
    best_map = -1.0
    best_path = run_dir / "best.pt"
    log(f"Training for epochs={train_epochs}, batch_size={batch_size}, eval_batch_size={eval_batch_size}")
    for epoch in range(1, train_epochs + 1):
        model.train()
        train_losses: List[float] = []
        for batch in tqdm(train_loader, desc=f"{args.run_name} epoch {epoch}/{train_epochs}"):
            targets = batch["targets"].to(device)
            logits = model(batch["images"], device=device)
            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            train_losses.append(float(loss.item()))

        val_metrics = run_eval(model=model, loader=val_loader, device=device, pos_weight=pos_weight)
        epoch_row = {
            "epoch": epoch,
            "train_loss": float(np.mean(train_losses)) if train_losses else 0.0,
            "val_loss": val_metrics["loss"],
            "val_map": val_metrics["map"],
            "val_macro_f1": val_metrics["macro_f1"],
            "val_micro_f1": val_metrics["micro_f1"],
        }
        history.append(epoch_row)
        log(
            f"Epoch {epoch}/{train_epochs} -> train_loss={epoch_row['train_loss']:.4f} "
            f"val_map={epoch_row['val_map']:.4f} val_macro_f1={epoch_row['val_macro_f1']:.4f}"
        )
        if val_metrics["map"] > best_map:
            best_map = val_metrics["map"]
            torch.save(
                {
                    "model_state": model.state_dict(),
                    "label_vocab": label_vocab,
                    "config": cfg,
                },
                best_path,
            )
            log(f"Saved new best checkpoint at epoch {epoch}: {best_path}")

    log("Running final test evaluation using best checkpoint...")
    checkpoint = torch.load(best_path, map_location=device)
    model.load_state_dict(checkpoint["model_state"])
    test_metrics = run_eval(model=model, loader=test_loader, device=device, pos_weight=pos_weight)
    payload = {
        "run_name": args.run_name,
        "mode": args.mode,
        "num_labels": num_labels,
        "best_val_map": best_map,
        "history": history,
        "test_metrics": test_metrics,
    }
    write_json(run_dir / "metrics.json", payload)
    write_json(run_dir / "label_vocab.json", label_vocab)
    with (run_dir / "overrides.json").open("w", encoding="utf-8") as handle:
        json.dump(parse_overrides_json(args.overrides_json), handle, indent=2, ensure_ascii=True)
    log(f"Training complete for {args.run_name}")
    print(json.dumps({"run_name": args.run_name, "test_map": test_metrics["map"]}, indent=2), flush=True)


if __name__ == "__main__":
    train()