Runners

Runners are used to execute inference or training of a model. They are initialised with a model, a dataset, callbacks and options. They are executed with the run function which takes no arguments.

Inference

mozuma.torch.runners.TorchInferenceRunner

Runner for inference tasks on PyTorch models

Supports CPU or single GPU inference.

Attributes:

Name	Type	Description
model	TorchModel	The PyTorch model to run inference
dataset	TorchDataset	Input dataset for the runner
callbacks	List[TorchRunnerCallbackType]	Callbacks to save features, labels or bounding boxes
options	TorchRunnerOptions	PyTorch options

run(self) -> None

Runs inference

mozuma.torch.runners.TorchInferenceMultiGPURunner

Runner for inference tasks on PyTorch models

Supports CPU and multi-GPU inference with native torch backends.

Attributes:

Name	Type	Description
model	TorchModel	The PyTorch model to run inference
dataset	TorchDataset	Input dataset for the runner
callbacks	List[TorchRunnerCallbackType]	Callbacks to save features, labels or bounding boxes
options	TorchMultiGPURunnerOptions	PyTorch multi-gpu options

run(self) -> None

Runs inference

Training

mozuma.torch.runners.TorchTrainingRunner dataclass

Runner for training tasks on PyTorch models

Supports CPU and multi-GPU training with multiple backends.

Attributes:

Name	Type	Description
model	TorchModel	The PyTorch model to run inference
dataset	Tuple[TorchTrainingDataset, TorchTrainingDataset]	Train and test datasets for the runner
callbacks	List[Union[BaseRunnerEndCallback, SaveModelState]]	Callback to save model weights plus one or more callbacks for when the runner ends.
options	TorchTrainingOptions	PyTorch training options

run(self) -> None

Runs training

Write your own runner

mozuma.runners.BaseRunner

A runner takes a model and run an action on it (inference, training...)

It must implement the run function and call the callbacks to save prediction or model weights.

Attributes:

Name	Type	Description
model	_ModelType	The model object to run the action against
dataset	_DataType	The input dataset
callbacks	List[_CallbackType]	Callbacks to save model state or predictions
options	_OptionsType	Options of the runner (devices...)

Note

The _ModelType, _DataType, _CallbackType and _OptionsType are generic types that should be specified when implementing a runner.

Examples:

This is an example of a runner that applies a function to each element of list dataset. It passes the returned data to the save_features callback.

from mozuma.callbacks.base import (
    BaseSaveFeaturesCallback,
    callbacks_caller
)
from mozuma.runners import BaseRunner

class NumpySumRunnerExample(BaseRunner[
    Callable,                   # _ModelType
    List,                       # _DataType
    BaseSaveFeaturesCallback,   # _CallbackType
    None                        # _OptionType
]):
    def run(self):
        for index, data in enumerate(self.dataset):
            # Helper function to call all matching callbacks
            callbacks_caller(
                self.callbacks,
                "save_features",
                index,
                self.model(data)
            )