QuarkGluonGAT — Checkpoint

3layer Graph Attention Network for quark/gluon jet classification trained on 139 306 calorimeter jet images (ECAL / HCAL / Tracks channels, 125×125 px).

Checkpoint format

{
    "model":     OrderedDict,   # load with model.load_state_dict()
    "optimizer": OrderedDict,
    "epoch":     int,
}

Loading example

import torch
from src.gnn import JetGNN

ckpt  = torch.load("gnn_checkpoint.pth", map_location="cpu")
model = JetGNN()
model.load_state_dict(ckpt["model"])
model.eval()

# model expects torch_geometric Data objects
# see src/dataset.py for the jet-image → graph conversion

Training

BCEWithLogitsLoss, Adam lr=1e-3, 20 epochs, batch size 32. Active pixels → nodes with features [x, y, ECAL, HCAL, Tracks]; k-NN graph (k=7).