208 lines
8.1 KiB
Python
208 lines
8.1 KiB
Python
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()
|