Skip to content

mlvbench.trainer

Trainer for fitting probes to pretrained vision models.

The ProbeTrainer is the core component that allows fitting task-specific probes to pretrained vision encoders. Creating a ProbeTrainer requires specifying task, probe type and data module. For example:

data_module = MyDataModule(...)
trainer = ProbeTrainer(
    task: "binary_segmentation",
    probe: "linear",
    data_module: data_module,
    batch_size: 256,
    max_steps: 1000,
)

The trainer then allows fitting and evaluating probes for various models:

model = build_model("...")
trainer.fit(model)

# Access the fitted probes via `model.probes`
print(f"Fitted {len(model.probes)} probes to the following layers:")
for probe in model.probes:
    print(probe.layer)

# Evaluate probes
test_loader = data_module.test_loader(batch_size=256)
trainer.evaluate(model, test_loader, "output/example/evaluation")

Have a look at the ProbeTrainer class for more options that allow customizing the training process.

ProbeTrainer

Trainer for fitting probes to pretrained vision models.

__init__(data_module, task, probe, batch_size, max_steps, optimizer='adamw', lr=0.001, weight_decay=0.0, early_stopping_patience=None, effective_batch_size=None, num_workers=0, val_every_n_steps=200, log_every_n_steps=50, enable_throughput_monitor=True, probe_types=None)

Initialize the trainer.

Parameters:

Name Type Description Default
data_module DataModule

The data module providing training and validation data.

 required
task Task | str

The task to train the probe on. Can be a string (name of the task) or a Task object.

 required
probe str

The name of the probe type to train (e.g. "linear"). Custom types must be registered via probe_types.

 required
batch_size int

The batch size for all data loaders.

 required
max_steps int

The maximum number of training steps.

 required
optimizer Literal['adamw', 'adam']

The optimizer to use. Can be "adamw" or "adam".

 'adamw'
lr float | list[float]

The learning rate(s) to use. If a list, different probes are trained for each learning rate and the best configuration per layer is selected based on validation loss.

 0.001
weight_decay float | list[float]

The weight decay value(s) to use. If a list, different probes are trained for each weight decay value and the best configuration per layer is selected based on the validation loss.

 0.0
early_stopping_patience int | None

If set, enable early stopping based on validation loss. The best probe state per layer is tracked and restored before the final evaluation, and training terminates once every layer has gone early_stopping_patience consecutive periodic validations without improvement.

 None
effective_batch_size int | None

The effective batch size for gradient accumulation. If None, defaults to batch_size (no accumulation). Must be a multiple of batch_size.

 None
num_workers int

Number of worker processes for the data loaders.

 0
val_every_n_steps int

Evaluate the probe every N training steps and after training completes. If None, the probe is evaluated only after training completes.

 200
log_every_n_steps int

Log the training progress every N training steps.

 50
enable_throughput_monitor bool

Whether to monitor throughput and memory usage during training and to write the recorded statistics to throughput.json. Disabling the throughput monitor avoids device synchronizations (see ThroughputMonitor).

 True
probe_types dict[str, type[Probe]] | None

Optional mapping of custom probe type names to Probe subclasses, merged with the built-in types.

 None

evaluate(model, dataloader, output_path, save_predictions=8)

Evaluate the attached probes on a single dataset subset.

This writes three files to the given output path:

  • results.parquet: Detailed results with one row per probe and sample.
  • summary.csv: Results averaged over the test set; one row per probe.
  • predictions.pt: A flat list of input images, targets and probe predictions. This allows visualizing results without having to restore the model.

Parameters:

Name Type Description Default
model Model

The model with attached probes to evaluate.

 required
dataloader Iterable[dict]

The data loader to evaluate on.

 required
output_path Path

The path to save the results to.

 required
save_predictions int

Maximum number of predictions to save in predictions.pt. Set to 0 to disable.

 8

fit(model, layers=None)

Fit probes for the given model.

Parameters:

Name Type Description Default
model Model

A fully set-up Model instance, already placed on its target device and optionally torch.compiled by the caller (see build_model). The trainer uses the model's device for training.

 required
layers list[str] | None

The model layer names to fit probes for. If None, all layers are used.

 None

Returns:

Type Description
list[dict[str, Any]]

The model with the trained probes attached (model.probes). The same model is also saved to output_path / "model.pt".