keras_backend

Classes:

Name	Description
KerasModelAdapter

Attributes:

Name	Type	Description
tf_version

tf_version module-attribute

tf_version = Version(tf.__version__)

KerasModelAdapter

KerasModelAdapter(model_description: AnyModelDescr, devices: Optional[Sequence[str]])

Bases: LocalModelAdapter[None, Any]

              flowchart TD
              bioimageio.core.backends.keras_backend.KerasModelAdapter[KerasModelAdapter]
              bioimageio.core._model_adapter.LocalModelAdapter[LocalModelAdapter]
              bioimageio.core._model_adapter.ModelAdapter[ModelAdapter]

                              bioimageio.core._model_adapter.LocalModelAdapter --> bioimageio.core.backends.keras_backend.KerasModelAdapter
                                bioimageio.core._model_adapter.ModelAdapter --> bioimageio.core._model_adapter.LocalModelAdapter
                



              click bioimageio.core.backends.keras_backend.KerasModelAdapter href "" "bioimageio.core.backends.keras_backend.KerasModelAdapter"
              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/_model_adapter.py

def __init__(
    self, model_description: AnyModelDescr, devices: Optional[Sequence[str]]
):
    super().__init__()
    self._model_descr = model_description
    self._input_ids = get_member_ids(model_description.inputs)
    self._output_ids = get_member_ids(model_description.outputs)
    self._input_axes = [
        tuple(a.id for a in get_axes_infos(t)) for t in model_description.inputs
    ]
    self._output_axes = [
        tuple(a.id for a in get_axes_infos(t)) for t in model_description.outputs
    ]
    if isinstance(model_description, v0_4.ModelDescr):
        self._input_is_optional = [False] * len(model_description.inputs)
    else:
        self._input_is_optional = [ipt.optional for ipt in model_description.inputs]

    self._devices = devices
    self.load()

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()