Skip to content

pytorch_backend ¤

Classes:

Name Description
PytorchModelAdapter

Functions:

Name Description
get_devices
load_torch_model
load_torch_state_dict

PytorchModelAdapter ¤

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

Bases: ModelAdapter


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

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


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

Methods:

Name Description
create

Creates model adapter based on the passed spec
forward

Run forward pass of model to get model predictions
load
unload

Unload model from any devices, freeing their memory.
Source code in src/bioimageio/core/backends/pytorch_backend.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    *,
    model_description: AnyModelDescr,
    devices: Optional[Sequence[Union[str, torch.device]]] = None,
    mode: Literal["eval", "train"] = "eval",
):
    super().__init__(model_description=model_description)
    weights = model_description.weights.pytorch_state_dict
    if weights is None:
        raise ValueError("No `pytorch_state_dict` weights found")

    devices = get_devices(devices)
    self._model = load_torch_model(weights, load_state=True, devices=devices)
    if mode == "eval":
        self._model = self._model.eval()
    elif mode == "train":
        self._model = self._model.train()
    else:
        assert_never(mode)

    self._mode: Literal["eval", "train"] = mode
    self._primary_device = devices[0]

create classmethod ¤

create(model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], *, devices: Optional[Sequence[str]] = None, weight_format_priority_order: Optional[Sequence[SupportedWeightsFormat]] = None)

Creates model adapter based on the passed spec Note: All specific adapters should happen inside this function to prevent different framework initializations interfering with each other

Used by:
Source code in src/bioimageio/core/backends/_model_adapter.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@final
@classmethod
def create(
    cls,
    model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr],
    *,
    devices: Optional[Sequence[str]] = None,
    weight_format_priority_order: Optional[Sequence[SupportedWeightsFormat]] = None,
):
    """
    Creates model adapter based on the passed spec
    Note: All specific adapters should happen inside this function to prevent different framework
    initializations interfering with each other
    """
    if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)):
        raise TypeError(
            f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}"
        )

    weights = model_description.weights
    errors: List[Exception] = []
    weight_format_priority_order = (
        DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER
        if weight_format_priority_order is None
        else weight_format_priority_order
    )
    # limit weight formats to the ones present
    weight_format_priority_order_present: Sequence[SupportedWeightsFormat] = [
        w for w in weight_format_priority_order if getattr(weights, w) is not None
    ]
    if not weight_format_priority_order_present:
        raise ValueError(
            f"None of the specified weight formats ({weight_format_priority_order}) is present ({weight_format_priority_order_present})"
        )

    for wf in weight_format_priority_order_present:
        if wf == "pytorch_state_dict":
            assert weights.pytorch_state_dict is not None
            try:
                from .pytorch_backend import PytorchModelAdapter

                return PytorchModelAdapter(
                    model_description=model_description, devices=devices
                )
            except Exception as e:
                errors.append(e)
        elif wf == "tensorflow_saved_model_bundle":
            assert weights.tensorflow_saved_model_bundle is not None
            try:
                from .tensorflow_backend import create_tf_model_adapter

                return create_tf_model_adapter(
                    model_description=model_description, devices=devices
                )
            except Exception as e:
                errors.append(e)
        elif wf == "onnx":
            assert weights.onnx is not None
            try:
                from .onnx_backend import ONNXModelAdapter

                return ONNXModelAdapter(
                    model_description=model_description, devices=devices
                )
            except Exception as e:
                errors.append(e)
        elif wf == "torchscript":
            assert weights.torchscript is not None
            try:
                from .torchscript_backend import TorchscriptModelAdapter

                return TorchscriptModelAdapter(
                    model_description=model_description, devices=devices
                )
            except Exception as e:
                errors.append(e)
        elif wf == "keras_hdf5":
            assert weights.keras_hdf5 is not None
            # keras can either be installed as a separate package or used as part of tensorflow
            # we try to first import the keras model adapter using the separate package and,
            # if it is not available, try to load the one using tf
            try:
                try:
                    from .keras_backend import KerasModelAdapter
                except Exception:
                    from .tensorflow_backend import KerasModelAdapter

                return KerasModelAdapter(
                    model_description=model_description, devices=devices
                )
            except Exception as e:
                errors.append(e)
        else:
            assert_never(wf)

    assert errors
    if len(weight_format_priority_order) == 1:
        assert len(errors) == 1
        raise errors[0]

    else:
        msg = (
            "None of the weight format specific model adapters could be created"
            + " in this environment."
        )
        raise ExceptionGroup(msg, errors)

forward ¤

forward(input_sample: Union[Sample, SampleBlock, SampleBlockWithOrigin]) -> Sample

Run forward pass of model to get model predictions

Note: sample id and stample stat attributes are passed through

Source code in src/bioimageio/core/backends/_model_adapter.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
def forward(
    self, input_sample: Union[Sample, SampleBlock, SampleBlockWithOrigin]
) -> Sample:
    """
    Run forward pass of model to get model predictions

    Note: sample id and stample stat attributes are passed through
    """
    unexpected = [mid for mid in input_sample.members if mid not in self._input_ids]
    if unexpected:
        warnings.warn(f"Got unexpected input tensor IDs: {unexpected}")

    input_arrays = [
        (
            None
            if (a := input_sample.members.get(in_id)) is None
            else a.transpose(in_order).data.data
        )
        for in_id, in_order in zip(self._input_ids, self._input_axes)
    ]
    output_arrays = self._forward_impl(input_arrays)
    assert len(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)
    ]
    return Sample(
        members={
            tid: out
            for tid, out in zip(
                self._output_ids,
                output_tensors,
            )
            if out is not None
        },
        stat=input_sample.stat,
        id=(
            input_sample.id
            if isinstance(input_sample, Sample)
            else input_sample.sample_id
        ),
    )

load ¤

load(*, devices: Optional[Sequence[str]] = None) -> None
Source code in src/bioimageio/core/backends/_model_adapter.py 
178
179
180
@final
def load(self, *, devices: Optional[Sequence[str]] = None) -> None:
    warnings.warn("Deprecated. ModelAdapter is loaded on initialization")

unload ¤

unload() -> None

Unload model from any devices, freeing their memory. The moder adapter should be considered unusable afterwards.

Source code in src/bioimageio/core/backends/pytorch_backend.py 
88
89
90
91
92
def unload(self) -> None:
    del self._model
    _ = gc.collect()  # deallocate memory
    assert torch is not None
    torch.cuda.empty_cache()  # release reserved memory

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 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def get_devices(
    devices: Optional[Sequence[Union[torch.device, str]]] = None,
) -> List[torch.device]:
    if not devices:
        torch_devices = [
            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        ]
    else:
        torch_devices = [torch.device(d) for d in devices]

    if len(torch_devices) > 1:
        warnings.warn(
            f"Multiple devices for single pytorch model not yet implemented; ignoring {torch_devices[1:]}"
        )
        torch_devices = torch_devices[:1]

    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 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
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,
            )
    return torch_model

load_torch_state_dict ¤

load_torch_state_dict(model: nn.Module, path: Union[Path, ZipPath, BytesReader], devices: Sequence[torch.device]) -> nn.Module
Source code in src/bioimageio/core/backends/pytorch_backend.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def load_torch_state_dict(
    model: nn.Module,
    path: Union[Path, ZipPath, BytesReader],
    devices: Sequence[torch.device],
) -> 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:
            state = torch.load(f, map_location=devices[0], weights_only=True)

    incompatible = model.load_state_dict(state)
    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