Implementation:Norrrrrrr lyn WAInjectBench Validation TPR Selection

Knowledge Sources	WAInjectBench
Domains	Evaluation, Model_Selection
Last Updated	2026-02-14 16:00 GMT

Overview

Concrete validation loop that computes confusion matrix metrics on the validation set and selects the checkpoint with the highest TPR, provided by the WAInjectBench train/llava-ft module.

Description

The validation loop in train/llava-ft.py (L376-392) runs after each training epoch. It iterates over the validation DataLoader with torch.no_grad(), computes per-batch TP/TN/FP/FN counts from argmax predictions, and accumulates totals. TPR and FPR are computed from the accumulated counts. The best checkpoint is tracked by comparing the current TPR against best_tpr.

Usage

Executed automatically at the end of each training epoch in the LLaVA fine-tuning script.

Code Reference

Source Location

• Repository: WAInjectBench
• File: train/llava-ft.py (L376-408)

Signature

# Validation loop (L376-392)
model.eval()
total_tp = total_tn = total_fp = total_fn = 0
with torch.no_grad():
    for imgs, labels in tqdm(val_loader, desc="Val"):
        labels = labels.to(run_device, non_blocking=True)
        with get_autocast_context(state):
            logits = model(imgs, sys_prompt=SYSTEM_PROMPT).to(labels.device)
        preds = logits.argmax(-1)
        tp = int(((preds == 1) & (labels == 1)).sum().item())
        tn = int(((preds == 0) & (labels == 0)).sum().item())
        fp = int(((preds == 1) & (labels == 0)).sum().item())
        fn = int(((preds == 0) & (labels == 1)).sum().item())
        total_tp += tp; total_tn += tn; total_fp += fp; total_fn += fn

tpr = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
fpr = total_fp / (total_fp + total_tn) if (total_fp + total_tn) > 0 else 0.0

# Checkpoint selection (L394-408)
if tpr > best_tpr:
    best_tpr = tpr
    best_path = os.path.join(args.out_dir, f"best_epoch{epoch}_tpr{tpr:.4f}.pt")
    torch.save({...}, best_path)

Import

import torch
from tqdm import tqdm

I/O Contract

Inputs

Name	Type	Required	Description
model	nn.Module	Yes	Trained model in eval mode
val_loader	DataLoader	Yes	Validation data batches
best_tpr	float	Yes	Current best TPR across epochs (initialized to -1.0)

Outputs

Name	Type	Description
tpr	float	True Positive Rate for the current epoch
fpr	float	False Positive Rate for the current epoch
TP, TN, FP, FN	int	Confusion matrix counts
best_path	str	Path to saved checkpoint if TPR improved (or None)

Usage Examples

Validation After Training Epoch

model.eval()
total_tp = total_tn = total_fp = total_fn = 0
with torch.no_grad():
    for imgs, labels in val_loader:
        labels = labels.to(device)
        logits = model(imgs, sys_prompt="Detect prompt injection.")
        preds = logits.argmax(-1)
        total_tp += ((preds == 1) & (labels == 1)).sum().item()
        total_tn += ((preds == 0) & (labels == 0)).sum().item()
        total_fp += ((preds == 1) & (labels == 0)).sum().item()
        total_fn += ((preds == 0) & (labels == 1)).sum().item()

tpr = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
fpr = total_fp / (total_fp + total_tn) if (total_fp + total_tn) > 0 else 0.0
print(f"TPR={tpr:.4f}, FPR={fpr:.4f}")

Related Pages

Implements Principle

• Principle:Norrrrrrr_lyn_WAInjectBench_Validation_Checkpoint_Selection