Add New Models#

The model wrapping class spora_bench.wrapper.SporaModelWrapper defines the methods expected by the pipelines:

Method	Benchmarks requiring an implementation
`compute_cell_tokens`	In-cohort cell phenotyping via linear probing
`embed_tissue`	Tissue-level tasks
`postprocess_tile_embeddings`	Tissue-level tasks (optional)
`predict_marker`	MPX-to-MPX virtual staining
`predict_instance_segmentation`	Cross-cohort cell instance segmentation
`predict_cell_types`	Cross-cohort cell-type annotation
`predict_markers_from_he`	H&E-to-MPX virtual staining
`predict_markers_from_he_batch`	H&E-to-MPX virtual staining (optional; override for cross-tile batched inference)

The API Reference documents the expected arguments and return types.

To implement a new model, create a class that inherits from SporaModelWrapper. The custom class must:

implement all interface methods required by the benchmarks you intend to run;
define an __init__ constructor that
- fully initializes the wrapped model,
- calls super().__init__(model_name=<model-name>), and
- accepts keyword arguments only; positional arguments are not supported.

See models/virtues/virtues_model.py, models/kronos/kronos_model.py, and models/resnet/resnet.py for examples.

import torch
from spora_bench.wrapper import SporaModelWrapper


class MyModelWrapper(SporaModelWrapper):
    def __init__(self, *, model_name: str, checkpoint_path: str):
        super().__init__(model_name=model_name)
        self.model = load_my_model(checkpoint_path)  # your own loader

    def compute_cell_tokens(self, dataset, tissue_id: str) -> torch.Tensor:
        tissue = dataset.get_tissue(tissue_id)
        ...  # return an (N_cells, D) tensor

Model Config#

Once the wrapper is implemented, write a model config. It is used to instantiate the model with hydra.utils.instantiate and must therefore follow this format:

model:
  _target_: path.to.model.MyModelClass
  model_name: my_model
  keywordarg1: <first-keyword-argument>
  keywordarg2: <second-keyword-argument>
  # ...

model_name also names the output directory output_dir/<model_name>/.