xie_x1andClaude Opus 5.5
0f9a3f496f
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>
251 lines
10 KiB
Python
251 lines
10 KiB
Python
#!/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()
|