111 lines
4.2 KiB
Python
111 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import random
|
|
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 supervised_alignment_loss
|
|
from losses.matching_loss import hungarian_matching_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 train(config: dict[str, Any], backbone_name: str, epochs: int | 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"
|
|
best_score = -1.0
|
|
epoch_count = epochs or config["training"]["epochs"]
|
|
|
|
for epoch in range(1, epoch_count + 1):
|
|
model.train()
|
|
running_loss = 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_facts"], config["training"]["valid_weight"]
|
|
)
|
|
text = hungarian_matching_loss(
|
|
outputs["text"], batch["text_facts"], config["training"]["valid_weight"]
|
|
)
|
|
alignment = supervised_alignment_loss(
|
|
outputs["visual"]["states"], batch["visual_facts"], visual.assignments
|
|
) + supervised_alignment_loss(
|
|
outputs["text"]["states"], batch["text_facts"], text.assignments
|
|
)
|
|
loss = visual.loss + text.loss + config["training"]["alignment_weight"] * alignment
|
|
loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(trainable, 1.0)
|
|
optimizer.step()
|
|
running_loss += float(loss.detach())
|
|
metrics = evaluate(model, validation_loader, device)
|
|
score = metrics["visual"]["f1"] + metrics["text"]["f1"]
|
|
summary = {
|
|
"epoch": epoch,
|
|
"loss": running_loss / max(1, len(train_loader)),
|
|
"metrics": metrics,
|
|
}
|
|
print(json.dumps(summary))
|
|
if score > best_score:
|
|
best_score = score
|
|
torch.save(
|
|
{"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, best_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"), default=None)
|
|
parser.add_argument("--epochs", type=int, 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))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|