SelectiveInferenceTask

Wrap a trained LightningModule to attach selection results during inference.

Usage

Source

SelectiveInferenceTask(
    task,
    score,
    acc_test_outputs=False,
    rc_metric=None,
)

The wrapper calls the wrapped model in inference mode and combines its predictions with selection outputs produced by a provided UncertaintyScore.

Examples

from seapig import SelectiveInferenceTask
from seapig.scores import EuclideanScore

score = EuclideanScore()
# score.fit(X=train_embeddings)  # fit before wrapping
selective_task = SelectiveInferenceTask(task=model, score=score)

Methods

Name Description
__init__() Create a SelectiveInferenceTask.
forward() Run the wrapped model and attach selection results.
get_risk_coverage_curve() Return the latest computed risk-coverage curve(s), or None if not available.
on_test_epoch_end() Log final computed test metrics once at the end of testing.
predict_step() Perform prediction and return predictions with selection outputs.
test_step() Execute a test step, updating metrics with selection information.

__init__()

Create a SelectiveInferenceTask.

Usage

Source

__init__(task, score, acc_test_outputs=False, rc_metric=None)
Parameters
task: LightningModule

A trained LightningModule that provides a extract(x) method returning a dict with predictions and any required fields for the score (e.g. embedding or `logit``). The task is deep-copied and set to eval mode to avoid side-effects during inference.

score: UncertaintyScore

An instance implementing select to compute a selection mask. It can be a generic UncertaintyScore or one of the concrete subclasses EmbeddingScore / LogitScore which operate on specific model outputs.

acc_test_outputs: bool = False

If True the per-batch outputs (predictions merged with selection results) are stored in self.test_outputs for later inspection.

rc_metric: RiskCoverageMetric | None = None
Optional metric to track risk-coverage during testing. If supplied it will be updated on each test step and the final curve can be retrieved via get_risk_coverage_curve().

forward()

Run the wrapped model and attach selection results.

Usage

Source

forward(batch)
Parameters
batch: Mapping[str, Any] | Sequence[Any] | torch.Tensor
Input batch passed directly to the task’s extract() method.
Returns
dict[str, torch.Tensor]
A dictionary containing the model’s predictions merged with the selection outputs produced by the configured UncertaintyScore. The selection dictionary always includes 'score' (the raw uncertainty values) and 'selected' (a boolean mask).

get_risk_coverage_curve()

Return the latest computed risk-coverage curve(s), or None if not available.

Usage

Source

get_risk_coverage_curve()

on_test_epoch_end()

Log final computed test metrics once at the end of testing.

Usage

Source

on_test_epoch_end()

predict_step()

Perform prediction and return predictions with selection outputs.

Usage

Source

predict_step(batch, batch_idx, dataloader_idx=0)

The wrapper calls self.forward and returns a mapping containing the model’s prediction together with the score and selected masks.

test_step()

Execute a test step, updating metrics with selection information.

Usage

Source

test_step(batch, batch_idx, dataloader_idx=0)
Parameters
batch: Mapping[str, Any] | Sequence[Any]

A batch from the test DataLoader. It is passed directly to the wrapped model’s extract method via self.forward.

batch_idx: int

Index of the current batch (required by Lightning but not used here).

dataloader_idx: int = 0
Index of the DataLoader when multiple loaders are used. Defaults to 0.
Notes

The method runs self.forward to obtain predictions together with the score and selected masks produced by the configured UncertaintyScore. It then updates any attached SelectiveMetric and RiskCoverageMetric (if supplied) and optionally records the full selection dict in self.test_outputs when acc_test_outputs=True. Metrics are logged via Lightning’s log_dict mechanism.