ConvProbeSegEvaluator¶

Registry key: "conv_probe_seg".

A natural step up from LinearProbeSegEvaluator: instead of a single 1×1 conv, the head is a small convolutional block that mixes channels and spatial context before the per-pixel classifier.

Sequential(
    Conv2d(D, hidden, kernel_size=1, bias=False),
    LayerNorm2d(hidden),
    Conv2d(hidden, hidden, kernel_size=3, padding=1, bias=False),
    LayerNorm2d(hidden),
)

followed by a 1×1 classifier and a bilinear upsample to the mask size. LayerNorm2d is channel-wise layer norm at each spatial position (canonical SAM / ViTDet impl) — works at any batch size, unlike BN.

Sits between LinearProbeSegEvaluator (~150 params) and UPerNetSegEvaluator (~9.8M params) on the head-capacity spectrum (~0.79M params at D=768, hidden_size=256).

Encoder + dataset contract¶

Identical to LinearProbeSegEvaluator:

Encoder must return Tensor[B, D, H', W'] (dense patch features). Every RadHarmony make_* recipe supports this via output_keys={"img", "mask"}:

from radharmony.evaluator.backbones import make_raddino
transform, encoder = make_raddino(device="cuda", output_keys={"img", "mask"})

Dataset must emit a "mask" key per sample (output_mask=True).

Usage¶

from radharmony.evaluator import ConvProbeSegEvaluator
from radharmony.evaluator.backbones import make_raddino
from radharmony.dataset import SIIMACRPTXDataset

transform, encoder = make_raddino(device="cuda", output_keys={"img", "mask"})

ds = SIIMACRPTXDataset(
    base_image_dir="/data/SIIM-ACR-PTX/",
    transform=transform,
    output_cls=True,
    output_mask=True,
    mask_output_dir="/scratch/ptx_masks/",
)

ev = ConvProbeSegEvaluator(
    encoder,
    dataset=ds,                  # k-fold mode
    num_classes=2,
    n_folds=5,
    epochs=20,
    lr=1e-3,
    conv_hidden_size=256,
    output_dir="outputs/conv_probe_seg",
)
df = ev.evaluate()
ev.save_results(df)

Constructor arguments¶

Specific to ConvProbeSegEvaluator:

Argument	Type	Default	Description
`conv_hidden_size`	`int`	`256`	Neck channel width — controls head capacity

Training arguments (shared with `LinearProbeSegEvaluator`)¶

Argument	Type	Default	Description
`n_folds`	`int`	`5`	Number of folds in k-fold mode; ignored in fixed-split mode
`n_train_samples`	`list[int]` \| `None`	`None`	Training-set size sweep; `None` = full train pool as a single point
`epochs`	`int`	`20`	Number of training epochs for the segmentation head
`lr`	`float`	`1e-3`	AdamW learning rate
`weight_decay`	`float`	`0.05`	AdamW weight decay
`class_weights`	`list[float]` \| `None`	`None`	Per-class weights for `nn.CrossEntropyLoss` — useful for class-imbalanced foreground (e.g. small lesions)
`save_predictions`	`bool`	`False`	Write per-image GT / Pred / Prob PNGs under `<output_dir>/{GT,Pred,Prob}/<run-id>/`
`early_stop_metric`	`str`	`"dice"`	Per-epoch val metric tracked for early stopping — any key in the seg metric panel; macro-averaged over the foreground classes
`early_stop_patience`	`int`	`5`	Stop training after this many epochs without improvement; best-epoch head state is restored
`val_fraction`	`float`	`0.1`	Fraction of the train subsample carved off as an inner val set for early stopping; `0.0` disables
`store_final_model`	`bool`	`False`	After `evaluate()`, also train a deployment head on all data → `final_head_`

Shared arguments (inherited from `BaseSegEvaluator`)¶

Argument	Type	Default	Description
`image_encoder`	`nn.Module` / callable	—	Frozen encoder; segmentation mode returns `Tensor[B, D, H, W]`
`dataset`	dataset	`None`	k-fold mode (mutually exclusive with `train_dataset`/`test_dataset`)
`train_dataset`	dataset	`None`	Fixed-split mode train pool
`test_dataset`	dataset	`None`	Fixed-split mode test set
`num_classes`	`int`	`2`	Number of segmentation classes, including background
`device`	`str`	`"cuda"`	`"cuda"`, `"cuda:N"`, or `"cpu"`
`batch_size`	`int`	`8`	DataLoader batch size
`num_workers`	`int`	`4`	DataLoader workers
`autocast_dtype`	`torch.dtype`	`torch.bfloat16`	Encoder-forward autocast dtype; `None` disables
`output_dir`	`str`	`None`	Directory for CSV output
`n_seeds`	`int`	`1`	Fixed-split only: train-subsample replicates
`base_seed`	`int`	`0`	RNG seed base
`n_bootstrap`	`int`	`0`	Fixed-split only: test-image bootstrap resamples
`bootstrap_seed`	`int`	`0`	RNG seed base for test-image resampling
`threshold_strategy`	`str`	`"youden"`	Threshold selection strategy for binary metrics

Output schema¶

Identical to LinearProbeSegEvaluator — same dice / iou / pixel_acc / tpr / tnr / ppv / npv metric panel, same per-row + summary CSV layout via save_results.

Notes¶

Same training loop, split modes, metrics, and prediction-dump behaviour as LinearProbeSegEvaluator; only _make_head is overridden.
The saving / reuse API (fit / save_head / load_head / predict, plus store_final_model) is inherited — see Saving and reuse. save_head records conv_hidden_size in the checkpoint so load_head rebuilds the matching conv-block head.
Encoder weights are frozen (requires_grad_(False)); only the conv-block head is updated. Optimizer: AdamW + CosineAnnealingLR.
bfloat16 autocast is used for the encoder forward pass; the head runs in fp32 for numerical stability.
For SIIM-ACR-PTX / Montgomery / Shenzhen workflows see notebooks/evaluator/fm_comparison_seg.ipynb and notebooks/evaluator/evaluator_backbone_verify_seg.ipynb.