def evaluate_model(
model_path: Path,
test_data_path: Path,
output_path: Path,
signal_len: int | None = None,
kmer_len: int | None = None,
batch_size: int = 512,
device: str = "cuda",
num_workers: int = 0,
emit_scores: Path | None = None,
mixed_precision: bool = False,
no_compile: bool = False,
) -> dict:
"""
Evaluate a trained model on test data.
Args:
model_path: Path to model checkpoint directory
test_data_path: Path to test dataset (.npz)
output_path: Path to save evaluation metrics (JSON)
signal_len: Signal length (if None, read from model config)
kmer_len: K-mer length (if None, read from model config)
batch_size: Batch size for evaluation
device: Device for inference
num_workers: DataLoader workers; 0 means auto (see
``resolve_dataloader_workers``) -- up to 8 on CUDA, 0 on CPU
emit_scores: If set, also write per-chunk scores to this .npz. A
confusion matrix is a summary at ONE threshold; the scores behind
it answer questions the summary cannot -- per-group error
breakdowns, paired model comparisons, calibration, and any
operating point other than the one that was reported. They are
computed either way, so this only decides whether they survive.
mixed_precision: Run the forward pass under autocast on CUDA. Off by
default so `eval test` matches `predict`, which never autocasts
(#264) -- opt in explicitly, same as `model train --mixed-precision`.
no_compile: Disable torch.compile even above the size threshold below.
Returns:
Dictionary with evaluation metrics
"""
logger.info(f"Loading model from {model_path}")
# Load model and config
model, config = load_model_from_checkpoint(model_path, device=device)
# Use config values if not provided
if signal_len is None:
signal_len = config["signal_len"]
if kmer_len is None:
kmer_len = config["kmer_len"]
model_type = config["model_name"]
num_out = config.get("num_out", 1)
left_context = config.get("left_context")
right_context = config.get("right_context")
seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK)
dwell_offset = config.get("dwell_offset", 0)
signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4)))
signal_mode = config.get("signal_mode", "both")
# GPU optimizations (forward-pass only, no training stability concerns)
if device != "cpu":
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision("high")
logger.info(f"Model: {model_type}")
logger.info(f"Signal length: {signal_len}")
logger.info(f"K-mer length: {kmer_len}")
logger.info(f"seq_encoding: {seq_encoding}, dwell_offset: {dwell_offset}")
# Load test dataset. Forward dwell_template_table from training config so
# the feature tensor has the same channel count as training — otherwise
# the feature-branch Conv1d crashes with a channel mismatch on models
# that were trained with per-AA dwell templates. Resolved against the
# model directory, not the CWD (issue #369) -- `train_model` copies the
# table there and records the bare filename.
logger.info(f"\nLoading test data from {test_data_path}")
dwell_template_table = resolve_dwell_template_path(config, base_dir=Path(model_path))
# Channel ablation (issue #358 follow-up): the checkpoint may have been
# trained on a SUBSET of this corpus's channels -- e.g. every arm of an
# ablation sweep scored against the same all-channel corpus. `None` when
# the checkpoint's declared list already matches the corpus (every
# pre-existing pair), so this is a no-op for anything that predates the
# feature. `feature_channel_names` (the corpus's own recorded list) must
# travel alongside the selection -- LeechDataset needs it to reconcile
# `selected_feature_channels` against the corpus's actual data by name.
corpus_feature_channels = corpus_recorded_feature_channels(test_data_path)
feature_channels_selection = resolve_corpus_feature_selection(
config.get("feature_channels", DEFAULT_FEATURE_CHANNELS), test_data_path
)
test_dataset = LeechDataset(
test_data_path,
signal_len=signal_len,
kmer_len=kmer_len,
model_type=model_type,
dwell_offset=dwell_offset,
left_context=left_context,
right_context=right_context,
seq_encoding=seq_encoding,
signal_kmer_context=signal_kmer_context,
# The encoding is the trained model's, not a preference: its sequence
# branch has a fixed channel count, so falling back to base_onehot
# cannot produce a working run — only a channel-count RuntimeError
# several steps later, naming neither the corpus nor the encoding (#230).
allow_encoding_fallback=False,
signal_mode=signal_mode,
dwell_template_table=dwell_template_table,
# The checkpoint's own recorded motif_offset -- guards a
# dwell_template_table whose content changed on disk since training
# against the position/motif_offset pairing it was trained on (#364).
motif_offset=config.get("motif_offset", 0),
feature_channel_names=corpus_feature_channels,
selected_feature_channels=feature_channels_selection,
)
# torch.compile the model (CUDA graph + kernel fusion). Same policy as
# `predict` (PR #253, threshold shared via TORCH_COMPILE_MIN_SAMPLES):
# below the threshold, compilation overhead (~15-30s) outweighs the
# speedup, so small runs auto-skip. Never on CPU -- a CPU graph buys
# nothing here and previously compiled unconditionally, mode=None
# included (#264).
if no_compile:
logger.info("torch.compile disabled (--no-compile)")
elif device == "cpu":
logger.info("torch.compile skipped (CPU)")
elif len(test_dataset) < TORCH_COMPILE_MIN_SAMPLES:
logger.info(
f"torch.compile auto-skipped ({len(test_dataset)} chunks < "
f"{TORCH_COMPILE_MIN_SAMPLES} threshold)"
)
elif hasattr(torch, "compile"):
try:
model = torch.compile(model, mode="reduce-overhead")
logger.info("torch.compile enabled (mode=reduce-overhead)")
except Exception as e:
logger.debug(f"torch.compile failed, using eager mode: {e}")
# Wrap model for unified forward pass
model_wrapper = ModelInferenceWrapper(model, model_type)
# Collate, the host-to-device copy and the forward pass run serially in
# whichever process owns the loader, so a worker-less loader on a GPU means
# one core feeding an accelerator that then waits (issue #205).
effective_workers = resolve_dataloader_workers(num_workers, device)
loader_kwargs: dict = {"num_workers": effective_workers}
if device != "cpu":
loader_kwargs["pin_memory"] = True
if effective_workers > 0:
loader_kwargs["prefetch_factor"] = 4
test_loader = DataLoader(
test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn, **loader_kwargs
)
logger.info(f"Test samples: {len(test_dataset)}")
# Run evaluation
logger.info("\nRunning evaluation...")
# One definition, shared with predict: num_out > 1 (issue #269). This used
# to be `> 2` here, so a 2-output cross-entropy model was scored as binary
# by `eval test` while `predict` tagged and calibrated it as multiclass.
is_multiclass = _is_multiclass_config(config)
# Same decision predict makes (never autocasts by default): opt in with
# --mixed-precision rather than autocasting on CUDA unconditionally (#264).
use_autocast = mixed_precision and device != "cpu"
if is_multiclass:
all_labels_list: list[np.ndarray] = []
all_preds_list: list[np.ndarray] = []
all_logits_list: list[np.ndarray] = []
model.eval()
with torch.inference_mode():
with Progress() as progress:
task = progress.add_task("[cyan]Evaluating...", total=len(test_loader))
for batch in test_loader:
with torch.amp.autocast("cuda", enabled=use_autocast):
logits = model_wrapper.forward_batch(batch, device)
# Cast once, upstream of every downstream use (top-k, ECE,
# calibration, saved scores). autocast already promotes
# softmax/log_softmax to fp32 on its own, but argmax and
# the archived logits below don't go through softmax --
# this is the one place that protects all of them
# uniformly, matching validate()'s pre-cast in training.py.
logits = logits.float()
preds = torch.argmax(logits, dim=-1).cpu().numpy()
all_preds_list.append(preds.reshape(-1))
# Labels never need the device round trip -- nothing here
# computes with them on GPU (#264).
all_labels_list.append(batch["label"].numpy().reshape(-1))
all_logits_list.append(logits.cpu().numpy())
progress.update(task, advance=1)
all_labels_arr = (
np.concatenate(all_labels_list) if all_labels_list else np.zeros(0, dtype=np.int64)
)
all_preds_arr = (
np.concatenate(all_preds_list) if all_preds_list else np.zeros(0, dtype=np.int64)
)
accuracy = accuracy_score(all_labels_arr, all_preds_arr)
macro_f1 = sk_f1_score(all_labels_arr, all_preds_arr, average="macro", zero_division=0.0)
weighted_f1 = sk_f1_score(
all_labels_arr, all_preds_arr, average="weighted", zero_division=0.0
)
cm = sk_confusion_matrix(all_labels_arr, all_preds_arr).tolist()
# Per-class report
report = classification_report(
all_labels_arr, all_preds_arr, output_dict=True, zero_division=0.0
)
# Top-K accuracy (K=3, 5)
all_logits_combined = np.concatenate(all_logits_list, axis=0)
top3_acc = _topk_accuracy(all_logits_combined, all_labels_arr, k=3)
top5_acc = _topk_accuracy(all_logits_combined, all_labels_arr, k=5)
# ECE (top-1 confidence)
logits_t = torch.from_numpy(all_logits_combined).float()
labels_t = torch.from_numpy(all_labels_arr).long()
probs_t = torch.softmax(logits_t, dim=-1)
ece = _expected_calibration_error_multiclass(probs_t, labels_t)
logger.info(f"ECE (uncalibrated): {ece:.4f}")
# Calibrated ECE if calibration params exist
ece_calibrated = None
from leech.calibration import apply_calibration, load_calibration
cal_params = load_calibration(model_path)
if cal_params is not None:
cal_logits = apply_calibration(logits_t, cal_params)
cal_probs_t = torch.softmax(cal_logits, dim=-1)
ece_calibrated = _expected_calibration_error_multiclass(cal_probs_t, labels_t)
cal_method = cal_params.get("method", "unknown")
logger.info(f"ECE (calibrated, method={cal_method}): {ece_calibrated:.4f}")
metrics: dict = {
"accuracy": accuracy,
"macro_f1": macro_f1,
"weighted_f1": weighted_f1,
"top3_accuracy": top3_acc,
"top5_accuracy": top5_acc,
"ece": ece,
"confusion_matrix": cm,
"classification_report": report,
"num_classes": num_out,
}
if ece_calibrated is not None:
metrics["ece_calibrated"] = ece_calibrated
# Stratified by junction disruption (issue #282), when the test
# corpus carries the field. `junction_indel == 0` (exact) vs `!= 0`
# (disrupted) is exactly the split the charging model's honest FPR
# ceiling is measured on.
junction_indels = _read_junction_indels(test_data_path, len(all_labels_arr))
if junction_indels is not None:
metrics["junction_stratified"] = _junction_stratified_multiclass_metrics(
junction_indels, all_labels_arr, all_preds_arr
)
for name, stratum in metrics["junction_stratified"].items():
if stratum["n"]:
logger.info(
f"Junction {name}: n={stratum['n']}, "
f"accuracy={stratum['accuracy']:.4f}, macro_f1={stratum['macro_f1']:.4f}"
)
# Cross-run vs within-run, when the split recorded which runs trained.
cross_run = _read_cross_run_mask(test_data_path, len(all_labels_arr))
if cross_run is not None:
metrics["run_stratified"] = _run_stratified_metrics(
cross_run, all_labels_arr, all_preds_arr, None
)
_log_run_strata(metrics["run_stratified"])
# Add metadata
metrics["model_path"] = str(model_path)
metrics["test_data_path"] = str(test_data_path)
metrics["num_samples"] = len(all_labels_arr)
metrics["model_type"] = model_type
# Which precision regime produced these numbers -- now opt-in and
# off by default (#264), so this is what makes a metrics.json from
# before that default flip distinguishable from one after it.
metrics["mixed_precision"] = use_autocast
# Print summary
logger.info(f"Accuracy: {accuracy:.4f}")
logger.info(f"Macro F1: {macro_f1:.4f}")
logger.info(f"Weighted F1: {weighted_f1:.4f}")
logger.info(f"Top-3 accuracy: {top3_acc:.4f}")
logger.info(f"Top-5 accuracy: {top5_acc:.4f}")
if emit_scores is not None:
_save_scores(
emit_scores,
test_data_path,
all_labels_arr,
torch.softmax(logits_t, dim=-1).numpy(),
)
save_metrics(metrics, output_path)
return metrics
# Binary evaluation path
all_labels_list: list[np.ndarray] = []
all_probs_list: list[np.ndarray] = []
model.eval()
with torch.inference_mode():
with Progress() as progress:
task = progress.add_task("[cyan]Evaluating...", total=len(test_loader))
for batch in test_loader:
# Forward pass (wrapper handles moving tensors and calling model correctly)
with torch.amp.autocast("cuda", enabled=use_autocast):
logits = model_wrapper.forward_batch(batch, device)
# Under autocast the final Linear emits fp16, and unlike
# softmax (which autocast promotes to fp32 on its own),
# sigmoid is not on that list -- computing it on the raw
# logits would round the saved probability and quantize the
# AUROC ranking (#264). Cast once, before either branch below.
logits = logits.float()
# Get predictions
if num_out > 1:
probs = torch.softmax(logits, dim=-1)[:, 1].cpu().numpy()
else:
probs = torch.sigmoid(logits).cpu().numpy()
# Labels never need the device round trip -- nothing here
# computes with them on GPU (#264).
all_labels_list.append(batch["label"].numpy().reshape(-1))
all_probs_list.append(probs.reshape(-1))
progress.update(task, advance=1)
# Compute metrics
all_labels_arr = (
np.concatenate(all_labels_list) if all_labels_list else np.zeros(0, dtype=np.float32)
)
all_probs_arr = (
np.concatenate(all_probs_list) if all_probs_list else np.zeros(0, dtype=np.float32)
)
all_preds_arr = (all_probs_arr > 0.5).astype(int)
# checkpoint_metric (config.json, from training) is threaded through so a
# run selected on a parametric metric (tpr_at_fpr:<f> /
# callable_at_precision:<p>, issue #280) reports that same metric here --
# this is "test_metrics.json" (the --output file): the test half of
# "reports them ... so runs can be compared after the fact". A run that
# used a plain metric (or predates checkpoint_metric) adds nothing new.
metrics = compute_metrics(
all_labels_arr,
all_preds_arr,
all_probs_arr,
checkpoint_metric=config.get("checkpoint_metric"),
)
# Add metadata
metrics["model_path"] = str(model_path)
metrics["test_data_path"] = str(test_data_path)
metrics["num_samples"] = len(all_labels_arr)
metrics["model_type"] = model_type
# Which precision regime produced these numbers -- now opt-in and off by
# default (#264), so this is what makes a metrics.json from before that
# default flip distinguishable from one after it.
metrics["mixed_precision"] = use_autocast
# Stratified by junction disruption (issue #282), when the test corpus
# carries the field -- see the multiclass branch's identical comment.
junction_indels = _read_junction_indels(test_data_path, len(all_labels_arr))
if junction_indels is not None:
metrics["junction_stratified"] = _junction_stratified_binary_metrics(
junction_indels, all_labels_arr, all_preds_arr, all_probs_arr
)
for name, stratum in metrics["junction_stratified"].items():
if stratum["n"]:
logger.info(
f"Junction {name}: n={stratum['n']}, "
f"accuracy={stratum.get('accuracy', float('nan')):.4f}"
)
# Cross-run vs within-run, when the split recorded which runs trained --
# a run-held-out test split is the only honest cross-run number.
cross_run = _read_cross_run_mask(test_data_path, len(all_labels_arr))
if cross_run is not None:
metrics["run_stratified"] = _run_stratified_metrics(
cross_run, all_labels_arr, all_preds_arr, all_probs_arr
)
_log_run_strata(metrics["run_stratified"])
if emit_scores is not None:
_save_scores(emit_scores, test_data_path, all_labels_arr, all_probs_arr)
# Print and save
print_metrics(metrics)
save_metrics(metrics, output_path)
return metrics