MLXID/GridBased
0
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: grid-based multi-photon reconstruction

Browse Source 
Code, configs, charge LUTs, best.pth and evaluation records for each
training run, the knife-edge / Siemens-star / XRF notebooks with
outputs, and the overnight run reports. Epoch checkpoints, inference
results and HTML exports stay local (see .gitignore).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
xie_x1andClaude Opus 5.5 committed 2026-10-20 07:36:49 +02:00
commit 0f9a3f496f
222 files changed
+53227
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

scripts/03_train.py
+250
View File
@@ -0,0 +1,250 @@
#!/usr/bin/env python3
"""
WP3 -- train the grid-based model.
python scripts/03_train.py --config Configs/train_grid_12keV.yaml
python scripts/03_train.py --config ... --epochs 3 --tag smoke
Every epoch the validation pass reports both the loss and the *detection*
metrics (efficiency / purity / position RMS), because the loss alone says
nothing about the working point.
"""
import argparse
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from omegaconf import OmegaConf
from torch.utils.data import DataLoader
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))
from datasets import GridFrameDataset, collate_with_photons # noqa: E402
from decode import decode # noqa: E402
from losses import GridLoss # noqa: E402
from metrics import DetectionStats # noqa: E402
from models import build_model, valid_mask # noqa: E402
LOSS_PARTS = ("objectness", "offset", "energy")
def shard_files(folder, split, shard_range):
first, last = shard_range
return [os.path.join(folder, split, f"shard_{index:02d}.h5")
for index in range(int(first), int(last) + 1)]
def prepare_output_dir(conf, tag=None):
date = datetime.now().strftime("%y%m%d")
base = Path(conf.experiment.results_dir)
run_index = 0
while True:
name = (f"{date}_{conf.experiment.name}"
+ (f"_{tag}" if tag else "") + f"_{run_index:02d}")
output_dir = base / name
if not output_dir.exists():
break
run_index += 1
(output_dir / "Models").mkdir(parents=True)
(output_dir / "Plots").mkdir(parents=True)
OmegaConf.save(conf, output_dir / "config.yaml")
return output_dir
def run_epoch(model, loader, loss_function, mask, optimizer=None,
grad_scaler=None, use_amp=False):
is_training = optimizer is not None
model.train(is_training)
totals = dict(loss=0.0, objectness=0.0, offset=0.0, energy=0.0)
n_samples = 0
for batch in loader:
frames, objectness, offset, energy = [t.cuda(non_blocking=True)
for t in batch[:4]]
with torch.autocast("cuda", enabled=use_amp):
prediction = model(frames)
loss, parts = loss_function(prediction.float(), objectness, offset,
energy, mask)
if is_training:
optimizer.zero_grad(set_to_none=True)
if grad_scaler is not None:
grad_scaler.scale(loss).backward()
grad_scaler.step(optimizer)
grad_scaler.update()
else:
loss.backward()
optimizer.step()
batch_size = frames.shape[0]
totals["loss"] += loss.item() * batch_size
for part in LOSS_PARTS:
totals[part] += parts[part].item() * batch_size
n_samples += batch_size
return {key: value / max(n_samples, 1) for key, value in totals.items()}
@torch.no_grad()
def run_detection_metrics(model, loader, conf, halo, frame_size, use_amp=False):
model.eval()
stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
match_radius=float(conf.eval.match_radius),
halo=halo, frame_size=frame_size)
for frames, objectness, offset, energy, photons in loader:
frames = frames.cuda(non_blocking=True)
with torch.autocast("cuda", enabled=use_amp):
prediction = model(frames)
points = decode(prediction.float(),
threshold=float(conf.eval.threshold),
merge_distance=float(conf.eval.merge_distance),
halo=halo,
predict_energy=bool(conf.model.predict_energy),
energy_scale=float(conf.data.energy_scale))
stats.update([frame_points.numpy() for frame_points in points],
[frame_truth.numpy() for frame_truth in photons])
return stats.summary()
def plot_curves(history, output_dir):
epochs = np.arange(1, len(history["train_loss"]) + 1)
figure, axes = plt.subplots(1, 3, figsize=(15, 4.2))
axes[0].plot(epochs, history["train_loss"], label="train")
axes[0].plot(epochs, history["val_loss"], label="val")
axes[0].set_yscale("log")
axes[0].set_xlabel("epoch")
axes[0].set_ylabel("total loss")
axes[0].legend()
axes[0].grid(alpha=.3)
axes[1].plot(epochs, history["efficiency"], label="efficiency")
axes[1].plot(epochs, history["purity"], label="purity")
axes[1].set_xlabel("epoch")
axes[1].set_ylim(0, 1.02)
axes[1].legend()
axes[1].grid(alpha=.3)
axes[2].plot(epochs, history["rms_1d_um"], label="RMS 1D")
axes[2].plot(epochs, history["sigma68_1d_um"], label=r"$\sigma_{68}$")
axes[2].set_xlabel("epoch")
axes[2].set_ylabel(r"residual [$\mu$m]")
axes[2].set_yscale("log")
axes[2].legend()
axes[2].grid(alpha=.3)
figure.tight_layout()
figure.savefig(output_dir / "Plots" / "training_curves.png", dpi=130)
plt.close(figure)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True)
parser.add_argument("--epochs", type=int, default=None)
parser.add_argument("--tag", default=None)
args = parser.parse_args()
conf = OmegaConf.load(args.config)
if args.epochs:
conf.training.epochs = args.epochs
torch.manual_seed(conf.training.seed)
np.random.seed(conf.training.seed)
output_dir = prepare_output_dir(conf, args.tag)
print(f"[train] output -> {output_dir}")
dataset_kwargs = dict(input_scale=float(conf.data.input_scale),
energy_scale=float(conf.data.energy_scale))
train_dataset = GridFrameDataset(
shard_files(conf.data.folder, "train", conf.data.train_shards),
augment=bool(conf.data.augment), name="train",
max_frames=conf.data.max_train_frames, **dataset_kwargs)
val_dataset = GridFrameDataset(
shard_files(conf.data.folder, "val", conf.data.val_shards),
augment=False, name="val", return_photons=True,
max_frames=conf.data.max_val_frames, **dataset_kwargs)
n_workers = int(conf.data.num_workers)
train_loader = DataLoader(train_dataset, batch_size=int(conf.data.batch_size_train),
shuffle=True, num_workers=n_workers, pin_memory=True,
drop_last=True, persistent_workers=n_workers > 0)
val_loader = DataLoader(val_dataset, batch_size=int(conf.data.batch_size_val),
shuffle=False, num_workers=n_workers, pin_memory=True,
collate_fn=collate_with_photons,
persistent_workers=n_workers > 0)
model = build_model(conf).cuda()
halo = model.rf_radius
frame_size = train_dataset.frame_size
mask = valid_mask(frame_size, frame_size, halo, device="cuda")
n_parameters = sum(p.numel() for p in model.parameters())
print(f"[train] GridNet: {n_parameters} parameters, receptive-field radius "
f"R={halo} px (halo), valid region "
f"{frame_size-2*halo}x{frame_size-2*halo} of {frame_size}x{frame_size}")
loss_function = GridLoss(conf).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=float(conf.training.lr),
weight_decay=float(conf.training.weight_decay))
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, "min", factor=float(conf.training.scheduler_factor),
patience=int(conf.training.scheduler_patience))
use_amp = bool(conf.training.amp)
grad_scaler = torch.amp.GradScaler("cuda") if use_amp else None
history = {key: [] for key in ("train_loss", "val_loss", "efficiency",
"purity", "rms_1d_um", "sigma68_1d_um")}
best_loss = float("inf")
for epoch in range(1, int(conf.training.epochs) + 1):
started = time.time()
train_metrics = run_epoch(model, train_loader, loss_function, mask,
optimizer, grad_scaler, use_amp)
val_metrics = run_epoch(model, val_loader, loss_function, mask,
None, None, use_amp)
detection = run_detection_metrics(model, val_loader, conf, halo,
frame_size, use_amp)
scheduler.step(val_metrics["loss"])
history["train_loss"].append(train_metrics["loss"])
history["val_loss"].append(val_metrics["loss"])
history["efficiency"].append(detection["efficiency"])
history["purity"].append(detection["purity"])
history["rms_1d_um"].append(detection["rms_1d_um"])
history["sigma68_1d_um"].append(detection["sigma68_1d_um"])
print(f"[{epoch:3d}/{conf.training.epochs}] "
f"train {train_metrics['loss']:.4f} "
f"(obj {train_metrics['objectness']:.4f} "
f"off {train_metrics['offset']:.5f} "
f"ene {train_metrics['energy']:.5f}) "
f"val {val_metrics['loss']:.4f} | "
f"eff {detection['efficiency']*100:5.2f}% "
f"pur {detection['purity']*100:5.2f}% "
f"| RMS {detection['rms_1d_um']:.3f} um "
f"s68 {detection['sigma68_1d_um']:.3f} um "
f"| lr {optimizer.param_groups[0]['lr']:.1e} "
f"| {time.time()-started:.0f}s", flush=True)
checkpoint = dict(model=model.state_dict(),
conf=OmegaConf.to_container(conf),
epoch=epoch, metrics=detection)
if val_metrics["loss"] < best_loss:
best_loss = val_metrics["loss"]
torch.save(checkpoint, output_dir / "Models" / "best.pth")
if epoch in list(conf.training.checkpoint_epochs) or epoch == conf.training.epochs:
torch.save(checkpoint, output_dir / "Models" / f"E{epoch}.pth")
torch.save(checkpoint, output_dir / "Models" / "last.pth")
plot_curves(history, output_dir)
with open(output_dir / "history.json", "w") as handle:
json.dump(history, handle, indent=1)
print(f"[train] done. best val loss {best_loss:.4f}. "
f"models in {output_dir/'Models'}")
if __name__ == "__main__":
main()