Skip to content

mlvbench.tasks

Tasks.

Tasks allow reusing the trainer for different probing experiments by specifying

  • How to obtain groung truth for a batch.
  • The loss function.
  • How to evaluate a probe.

build_task(name)

Build the task with the given name.

Parameters:

Name Type Description Default
name str

The name of the task.

 required

Returns:

Type Description
Task

The task object.

Task interface.

Evaluator

Bases: ABC

Base class for evaluators.

An evaluator is used to compute and aggregate performance metrics over a dataset (subset). The following simplified code shows how the evaluator is used by the trainer:

evaluator = task.get_evaluator()

for batch in dataset:
    ...
    evaluator.update(batch, prediction, target, weight)

summary, results, predictions = evaluator.finalize()

finalize() abstractmethod

Finalize the evaluation and return the aggregated results.

Returns:

Type Description
tuple[dict[str, float], DataFrame, list[dict[str, Any]]]

A 3-tuple of (summary, results, predictions). summary holds the aggregated metrics for the whole subset, results holds one row per sample, and predictions holds one record per stored sample as returned by collect_predictions. predictions may be an empty list.

update(batch, prediction, target, weight=None) abstractmethod

Update the evaluator with data from a single batch.

Task

Bases: ABC

Base class for tasks.

A task specifies the prediction target, the loss function and the evaluator.

configure(model)

Configure the task for the given model.

Called by the trainer before any data preparation or probe fitting. Override to store model attributes (e.g., patch size) that are needed during target preparation or evaluation.

Parameters:

Name Type Description Default
model Model

The model whose features will be probed.

 required

criterion(prediction, target, weight=None) abstractmethod

Loss function for the task.

Parameters:

Name Type Description Default
prediction Tensor

The prediction from the probe, a float32 tensor whose shape depends on the task.

 required
target Tensor

The target for the task, a float32 tensor whose shape depends on the task.

 required
weight Tensor | None

Optional per-element float32 weight tensor with the same shape as target.

 None

Returns:

Type Description
Tensor

The scalar loss for the task, a float32 tensor of shape ().

get_evaluator(save_predictions=8) abstractmethod

Return an evaluator for the task.

Parameters:

Name Type Description Default
save_predictions int

The maximum number of samples for which the evaluator keeps raw predictions. Set to 0 to disable.

 8

prepare_target(batch) abstractmethod

Prepare the target for the probe.

Parameters:

Name Type Description Default
batch dict[str, Tensor]

A batch from the dataset.

 required

Returns:

Type Description
tuple[Tensor, Tensor | None]

A 2-tuple of (target, weight). target is a float32 tensor whose shape depends on the task (e.g. (B, N, 1)). weight is None if all samples should be weighted equally, otherwise a float32 tensor with the same shape as target holding per-element weights.