pytorch_backend

Classes:

Name	Description
PytorchModelAdapter
TorchNNModuleLike

Functions:

Name	Description
get_devices
load_torch_model
load_torch_state_dict

PytorchModelAdapter

PytorchModelAdapter(model_description: AnyModelDescr, mode: Literal['eval', 'train'] = 'eval', devices: Optional[Sequence[str]] = None)

Bases: LocalModelAdapter[torch.device, nn.Module]

              flowchart TD
              bioimageio.core.backends.pytorch_backend.PytorchModelAdapter[PytorchModelAdapter]
              bioimageio.core._model_adapter.LocalModelAdapter[LocalModelAdapter]
              bioimageio.core._model_adapter.ModelAdapter[ModelAdapter]

                              bioimageio.core._model_adapter.LocalModelAdapter --> bioimageio.core.backends.pytorch_backend.PytorchModelAdapter
                                bioimageio.core._model_adapter.ModelAdapter --> bioimageio.core._model_adapter.LocalModelAdapter
                



              click bioimageio.core.backends.pytorch_backend.PytorchModelAdapter href "" "bioimageio.core.backends.pytorch_backend.PytorchModelAdapter"
              click bioimageio.core._model_adapter.LocalModelAdapter href "" "bioimageio.core._model_adapter.LocalModelAdapter"
              click bioimageio.core._model_adapter.ModelAdapter href "" "bioimageio.core._model_adapter.ModelAdapter"
            

Methods:

Name	Description
close	Close the model adapter, freeing any resources.
forward	Run forward pass of model to get model predictions
load
unload	Unload model from any devices, freeing their memory.

Attributes:

Name	Type	Description
model_descr	AnyModelDescr

Source code in src/bioimageio/core/backends/pytorch_backend.py

def __init__(
    self,
    model_description: AnyModelDescr,
    mode: Literal["eval", "train"] = "eval",
    devices: Optional[Sequence[str]] = None,
):
    weights = model_description.weights.pytorch_state_dict
    if weights is None:
        raise ValueError("No `pytorch_state_dict` weights found")

    self._weights = weights
    self._mode: Literal["eval", "train"] = mode
    super().__init__(model_description=model_description, devices=devices)

model_descr property

model_descr: AnyModelDescr

close

close()

Close the model adapter, freeing any resources.

Note

The moder adapter should be considered unusable afterwards.

Source code in src/bioimageio/core/_model_adapter.py

def close(self):
    """Close the model adapter, freeing any resources.

    Note:
        The moder adapter should be considered unusable afterwards.
    """
    self.unload()

forward

forward(inputs: PerMember[Optional[Tensor]]) -> PerMember[Optional[Tensor]]

Run forward pass of model to get model predictions

Note: sample id and stample stat attributes are passed through

Parameters:

Name	Type	Description	Default
inputs	PerMember[Optional[Tensor]]	input tensors for the model, keyed by member id	required

Source code in src/bioimageio/core/_model_adapter.py

def forward(
    self, inputs: PerMember[Optional[Tensor]]
) -> PerMember[Optional[Tensor]]:
    """
    Run forward pass of model to get model predictions

    Note: sample id and stample stat attributes are passed through

    Args:
        inputs: input tensors for the model, keyed by member id
    """
    if not self._loaded:
        raise RuntimeError("Model must be `.load()`ed before calling forward()")

    unexpected = [mid for mid in inputs if mid not in self._input_ids]
    if unexpected:
        warnings.warn(f"Got unexpected input tensor IDs: {unexpected}")

    input_arrays = [
        (
            None
            if (a := inputs.get(in_id)) is None
            else a.transpose(in_order).to_numpy()
        )
        for in_id, in_order in zip(self._input_ids, self._input_axes)
    ]

    logger.debug(
        "NN input shapes: {}",
        [a.shape if a is not None else None for a in input_arrays],
    )
    device, model = self._model_queue.get()
    try:
        output_arrays = self._forward_impl(device, model, input_arrays)
    finally:
        self._model_queue.put((device, model))

    logger.debug(
        "NN output shapes: {}",
        [a.shape if a is not None else None for a in output_arrays],
    )
    if len(output_arrays) > len(self._output_ids):
        warnings.warn(
            f"Model produced more outputs ({len(output_arrays)}) than specified in the model description ({len(self._output_ids)}). Extra outputs will be ignored."
        )
        output_arrays = output_arrays[: len(self._output_ids)]

    output_tensors = [
        None if a is None else Tensor(a, dims=d)
        for a, d in zip(output_arrays, self._output_axes)
    ]
    outputs = {
        tid: out
        for tid, out in zip(
            self._output_ids,
            output_tensors,
        )
        if out is not None
    }
    outputs = restore_batch_multi_index(inputs, outputs)
    return outputs

load

load() -> None

Source code in src/bioimageio/core/_model_adapter.py

def load(self) -> None:
    devices = self._devices
    self._model_queue: LifoQueue[Tuple[DeviceType, ModelType]] = LifoQueue()
    parsed_devices = self._parse_devices(devices)
    assert parsed_devices
    # prioritize devices by order specified by user
    device_exceptions: Dict[str, Exception] = {}
    self._initialized_devices: List[str] = []
    for d in parsed_devices[::-1]:
        try:
            model = self._init_model_on_device(d)
        except Exception as e:
            device_exceptions[str(d)] = e
        else:
            self._model_queue.put((d, model))
            self._initialized_devices.insert(0, str(d))

    if self._model_queue.empty():
        if len(device_exceptions) == 1:
            raise next(iter(device_exceptions.values()))
        else:
            raise ExceptionGroup(
                "Failed to initialize model on any of the requested devices.",
                list(device_exceptions.values())[::-1],
            )

    if device_exceptions:
        logger.warning(
            "Failed to initialize model on some of the requested devices. Successfully initialized on {}, but got the following errors for other devices: {}",
            self._initialized_devices,
            device_exceptions,
        )

    super().load()

unload

unload()

Unload model from any devices, freeing their memory.

Note

The moder adapter should be considered unusable afterwards.

Source code in src/bioimageio/core/_model_adapter.py

def unload(self):
    for _ in range(len(self._initialized_devices)):
        device, model = self._model_queue.get()
        try:
            self._cleanup_pre_model_deletion(device, model)
        except Exception as e:
            logger.warning(
                "Got error during pre-deletion cleanup on device {}: {}", device, e
            )
        finally:
            del model
        try:
            self._cleanup_post_model_deletion(device)
        except Exception as e:
            logger.warning(
                "Got error during post-deletion cleanup on device {}: {}", device, e
            )

    _ = gc.collect()  # deallocate memory
    super().unload()

TorchNNModuleLike

Bases: Protocol

              flowchart TD
              bioimageio.core.backends.pytorch_backend.TorchNNModuleLike[TorchNNModuleLike]

              

              click bioimageio.core.backends.pytorch_backend.TorchNNModuleLike href "" "bioimageio.core.backends.pytorch_backend.TorchNNModuleLike"
            

Methods:

Name	Description
eval	Set model to eval mode
forward
load_state_dict
to

eval

eval() -> Self

Set model to eval mode

Source code in src/bioimageio/core/backends/pytorch_backend.py

def eval(self) -> Self:
    """Set model to eval mode"""
    return self

forward abstractmethod

forward(*input: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], List[torch.Tensor]]

Source code in src/bioimageio/core/backends/pytorch_backend.py

@abstractmethod
def forward(
    self, *input: torch.Tensor
) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], List[torch.Tensor]]: ...

load_state_dict abstractmethod

load_state_dict(state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False) -> Self

Source code in src/bioimageio/core/backends/pytorch_backend.py

@abstractmethod
def load_state_dict(
    self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False
) -> Self: ...

to abstractmethod

to(*, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, non_blocking: bool = False) -> Self

Source code in src/bioimageio/core/backends/pytorch_backend.py

@abstractmethod
def to(
    self,
    *,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    non_blocking: bool = False,
) -> Self: ...

get_devices

get_devices(devices: Optional[Sequence[Union[torch.device, str]]] = None) -> List[torch.device]

Source code in src/bioimageio/core/backends/pytorch_backend.py

def get_devices(
    devices: Optional[Sequence[Union[torch.device, str]]] = None,
) -> List[torch.device]:
    if not devices:
        if torch.cuda.is_available():
            torch_devices = [
                torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())
            ]
        elif torch.backends.mps.is_available():
            torch_devices = [torch.device("mps")]
        else:
            try:
                if (
                    torch.accelerator.is_available()
                    and (current_accelerator := torch.accelerator.current_accelerator())
                    is not None
                ):
                    torch_devices = [current_accelerator]
                else:
                    torch_devices = [torch.device("cpu")]
            except Exception:
                torch_devices = [torch.device("cpu")]
    else:
        torch_devices = [torch.device(d) for d in devices]

    return torch_devices

load_torch_model

load_torch_model(weight_spec: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr], *, load_state: bool = True, devices: Optional[Sequence[Union[str, torch.device]]] = None) -> nn.Module

Source code in src/bioimageio/core/backends/pytorch_backend.py

def load_torch_model(
    weight_spec: Union[
        v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr
    ],
    *,
    load_state: bool = True,
    devices: Optional[Sequence[Union[str, torch.device]]] = None,
) -> nn.Module:
    custom_callable = import_callable(
        weight_spec.architecture,
        sha256=(
            weight_spec.architecture_sha256
            if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr)
            else weight_spec.sha256
        ),
    )
    model_kwargs = (
        weight_spec.kwargs
        if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr)
        else weight_spec.architecture.kwargs
    )
    torch_model = custom_callable(**model_kwargs)

    if not isinstance(torch_model, nn.Module):
        if isinstance(
            weight_spec.architecture,
            (v0_4.CallableFromFile, v0_4.CallableFromDepencency),
        ):
            callable_name = weight_spec.architecture.callable_name
        else:
            callable_name = weight_spec.architecture.callable

        raise ValueError(f"Calling {callable_name} did not return a torch.nn.Module.")

    if load_state or devices:
        use_devices = get_devices(devices)
        torch_model = torch_model.to(use_devices[0])
        if load_state:
            torch_model = load_torch_state_dict(
                torch_model,
                path=download(weight_spec),
                devices=use_devices,
                strict=weight_spec.strict
                if isinstance(weight_spec, v0_5.PytorchStateDictWeightsDescr)
                else True,
            )
    return torch_model

load_torch_state_dict

load_torch_state_dict(model: nn.Module, path: Union[Path, ZipPath, BytesReader], devices: Sequence[torch.device], strict: bool = True) -> nn.Module

Source code in src/bioimageio/core/backends/pytorch_backend.py

def load_torch_state_dict(
    model: nn.Module,
    path: Union[Path, ZipPath, BytesReader],
    devices: Sequence[torch.device],
    strict: bool = True,
) -> nn.Module:
    model = model.to(devices[0])
    if isinstance(path, (Path, ZipPath)):
        ctxt = path.open("rb")
    else:
        ctxt = nullcontext(BytesIO(path.read()))

    with ctxt as f:
        assert not isinstance(f, TextIOWrapper)
        if Version(str(torch.__version__)) < Version("1.13"):
            state = torch.load(f, map_location=devices[0])
        else:
            try:
                state = torch.load(f, map_location=devices[0], weights_only=True)
            except Exception as e:
                msg = (
                    f"Failed to load weights with `weights_only=True`: {e}\n\n"
                    + "This usually means the weights file contains non-tensor objects"
                    + " (e.g. numpy arrays, custom classes, or nested dicts with"
                    + " metadata). The BioImage.IO spec requires a pure state dict —"
                    + " an OrderedDict mapping parameter names to tensors only.\n\n"
                    + "To fix this, extract only the state dict from your checkpoint:\n\n"
                    + "    import torch\n"
                    + "    checkpoint = torch.load('original.pth', weights_only=False)\n"
                    + "    # Inspect keys, e.g.: checkpoint.keys()"
                    + " -> dict_keys(['model', 'optimizer', ...])\n"
                    + "    torch.save(checkpoint['model'], 'weights.pt')\n\n"
                    + "Then reference 'weights.pt' in your bioimageio.yaml."
                )
                raise ValueError(msg) from e

    incompatible = model.load_state_dict(state, strict=strict)
    if (
        isinstance(incompatible, tuple)
        and hasattr(incompatible, "missing_keys")
        and hasattr(incompatible, "unexpected_keys")
    ):
        if incompatible.missing_keys:
            logger.warning("Missing state dict keys: {}", incompatible.missing_keys)

        if hasattr(incompatible, "unexpected_keys") and incompatible.unexpected_keys:
            logger.warning(
                "Unexpected state dict keys: {}", incompatible.unexpected_keys
            )
    else:
        logger.warning(
            "`model.load_state_dict()` unexpectedly returned: {} "
            + "(expected named tuple with `missing_keys` and `unexpected_keys` attributes)",
            (s[:20] + "..." if len(s := str(incompatible)) > 20 else s),
        )

    return model