Update to new pytorch_state_dict weight spec

Author: constantinpape

bioimageio/core/prediction_pipeline/_model_adapters/_keras_model_adapter.py

@@ -1,6 +1,7 @@
 import warnings
 from typing import List, Optional, Sequence
 
+# by default, we use the keras integrated with tensorflow
 try:
     from tensorflow import keras
 except Exception:

bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py

@@ -52,8 +52,10 @@ def _unload(self) -> None:
 
     @staticmethod
     def get_nn_instance(model_node: nodes.Model, **kwargs):
-        assert isinstance(model_node.source, nodes.ImportedSource)
-
-        joined_kwargs = {} if model_node.kwargs is missing else dict(model_node.kwargs)
+        weight_spec = model_node.weights.get("pytorch_state_dict")
+        assert weight_spec is not None
+        assert isinstance(weight_spec.architecture, nodes.ImportedSource)
+        model_kwargs = weight_spec.kwargs
+        joined_kwargs = {} if model_kwargs is missing else dict(model_kwargs)
         joined_kwargs.update(kwargs)
-        return model_node.source(**joined_kwargs)
+        return weight_spec.architecture(**joined_kwargs)

bioimageio/core/weight_converter/torch/utils.py

@@ -1,19 +1,11 @@
 import torch
-from marshmallow import missing
-
-
-# NOTE: copied from tiktorch; this should go into python-bioimageio and then we use it from there
-def get_nn_instance(node, **kwargs):
-    joined_kwargs = {} if node.kwargs is missing else dict(node.kwargs)  # type: ignore
-    joined_kwargs.update(kwargs)
-    model = node.source(**joined_kwargs)
-    return model
+from bioimageio.core.prediction_pipeline._model_adapters._pytorch_model_adapter import PytorchModelAdapter
 
 
 # additional convenience for pytorch state dict, eventually we want this in python-bioimageio too
 # and for each weight format
 def load_model(node):
-    model = get_nn_instance(node)
+    model = PytorchModelAdapter.get_nn_instance(node)
     state = torch.load(node.weights["pytorch_state_dict"].source)
     model.load_state_dict(state)
     model.eval()