from __future__ import annotations import argparse import json import random import time from pathlib import Path from typing import Any import numpy as np import torch import yaml from torch.utils.data import DataLoader from data.dataset import EvidenceDataset, collate_evidence from eval.evaluate_slots import evaluate from losses.alignment_loss import cross_modal_alignment_loss from losses.matching_loss import hungarian_matching_loss from losses.prototype_loss import canonical_prototype_loss from models.evidence_model import EvidenceExtractorModel def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def trainable_state_dict(model: EvidenceExtractorModel) -> dict[str, torch.Tensor]: names = {name for name, parameter in model.named_parameters() if parameter.requires_grad} return {name: value for name, value in model.state_dict().items() if name in names} def train( config: dict[str, Any], backbone_name: str, epochs: int | None = None, resume: Path | None = None, ) -> Path: set_seed(config["seed"]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") data_root = Path(config["data"]["root"]) train_dataset = EvidenceDataset(data_root / "train.jsonl", config["data"]["image_size"]) validation_dataset = EvidenceDataset( data_root / "validation.jsonl", config["data"]["image_size"] ) loader_options = { "batch_size": config["training"]["batch_size"], "num_workers": config["training"]["num_workers"], "collate_fn": collate_evidence, } train_loader = DataLoader(train_dataset, shuffle=True, **loader_options) validation_loader = DataLoader(validation_dataset, **loader_options) model = EvidenceExtractorModel(config, backbone_name=backbone_name).to(device) trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] optimizer = torch.optim.AdamW( trainable, lr=config["training"]["learning_rate"], weight_decay=config["training"]["weight_decay"], ) output_dir = Path(config["training"]["output_dir"]) if backbone_name != config["model"]["backbone"]: output_dir = output_dir.parent / backbone_name output_dir.mkdir(parents=True, exist_ok=True) best_path = output_dir / "best.pt" last_path = output_dir / "last.pt" details_path = output_dir / "training_details.jsonl" best_score = -1.0 epoch_count = epochs or config["training"]["epochs"] start_epoch = 1 resume_path = resume if resume_path is None and config["training"].get("auto_resume") and last_path.exists(): resume_path = last_path if resume_path is not None: checkpoint = torch.load(resume_path, map_location=device, weights_only=True) model.load_state_dict(checkpoint["model"], strict=False) optimizer.load_state_dict(checkpoint["optimizer"]) start_epoch = int(checkpoint["epoch"]) + 1 best_score = float(checkpoint.get("best_score", -1.0)) total_parameters = sum(parameter.numel() for parameter in model.parameters()) trainable_parameters = sum(parameter.numel() for parameter in trainable) run_details = { "event": "run_start", "device": str(device), "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else None, "backbone": backbone_name, "total_parameters": total_parameters, "trainable_parameters": trainable_parameters, "train_samples": len(train_dataset), "validation_samples": len(validation_dataset), "start_epoch": start_epoch, "target_epochs": epoch_count, "config": config, } details_mode = "a" if start_epoch > 1 else "w" with details_path.open(details_mode, encoding="utf-8") as handle: handle.write(json.dumps(run_details) + "\n") print(json.dumps(run_details)) for epoch in range(start_epoch, epoch_count + 1): epoch_started = time.perf_counter() model.train() running = { "total": 0.0, "visual": 0.0, "text": 0.0, "alignment": 0.0, "prototype": 0.0, } for batch in train_loader: optimizer.zero_grad(set_to_none=True) outputs = model(batch["images"].to(device), batch["rules_text"]) visual = hungarian_matching_loss( outputs["visual"], batch["visual_targets"], config["training"]["type_weight"] ) text = hungarian_matching_loss( outputs["text"], batch["text_targets"], config["training"]["type_weight"] ) alignment = cross_modal_alignment_loss( outputs["visual"]["states"], batch["visual_targets"], visual.assignments, outputs["text"]["states"], batch["text_targets"], text.assignments, config["training"]["contrastive_temperature"], ) loss = visual.loss + text.loss + config["training"]["alignment_weight"] * alignment prototype_loss = loss.new_zeros(()) if model.prototype_encoder is not None: prototype_loss = canonical_prototype_loss( outputs["visual"]["states"], batch["visual_targets"], visual.assignments, model.prototype_encoder, ) + canonical_prototype_loss( outputs["text"]["states"], batch["text_targets"], text.assignments, model.prototype_encoder, ) loss = loss + config["prototype"]["weight"] * prototype_loss loss.backward() torch.nn.utils.clip_grad_norm_(trainable, 1.0) optimizer.step() running["total"] += float(loss.detach()) running["visual"] += float(visual.loss.detach()) running["text"] += float(text.loss.detach()) running["alignment"] += float(alignment.detach()) running["prototype"] += float(prototype_loss.detach()) metrics = evaluate(model, validation_loader, device) score = metrics["visual_fact"]["f1"] + metrics["text_all"]["f1"] summary = { "event": "epoch_end", "epoch": epoch, "duration_seconds": time.perf_counter() - epoch_started, "learning_rate": optimizer.param_groups[0]["lr"], "loss": { key: value / max(1, len(train_loader)) for key, value in running.items() }, "metrics": metrics, } print(json.dumps(summary)) with details_path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(summary) + "\n") if score > best_score: best_score = score torch.save( { "model": trainable_state_dict(model), "epoch": epoch, "metrics": metrics, "backbone": backbone_name, }, best_path, ) torch.save( { "model": trainable_state_dict(model), "optimizer": optimizer.state_dict(), "epoch": epoch, "metrics": metrics, "best_score": best_score, "backbone": backbone_name, }, last_path, ) return best_path def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml")) parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None) parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--resume", type=Path, default=None) args = parser.parse_args() config = yaml.safe_load(args.config.read_text(encoding="utf-8")) backbone_name = args.backbone or config["model"]["backbone"] print(train(config, backbone_name, epochs=args.epochs, resume=args.resume)) if __name__ == "__main__": main()