Creates model adapter for model_descritption
Source code in src/bioimageio/core/backends/__init__.py
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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 | def create_model_adapter(
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 for `model_descritption`"""
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, None) 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, 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, 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, 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, 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, devices=devices)
except Exception as e:
errors.append(e)
elif wf == "keras_v3":
assert not isinstance(weights, v0_4.WeightsDescr), (
"keras_v3 weights not supported for v0.4 specs"
)
assert weights.keras_v3 is not None
try:
from .keras_backend import KerasModelAdapter
return KerasModelAdapter(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)
|