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()