Merge pull request #177 from bioimage-io/torchscript

Rename pytorch_script to torchscript

Author: constantinpape

.github/workflows/build.yml

@@ -80,6 +80,7 @@ jobs:
       uses: conda-incubator/setup-miniconda@v2
       with:
         auto-update-conda: true
+        # we need mamba to resolve environment-tf
         mamba-version: "*"
         channel-priority: strict
         activate-environment: bio-core-tf

bioimageio/core/__main__.py

@@ -249,10 +249,10 @@ def convert_torch_weights_to_torchscript(
         output_path: Path = typer.Argument(..., help="Where to save the torchscript weights."),
         use_tracing: bool = typer.Option(True, help="Whether to use torch.jit tracing or scripting."),
     ) -> int:
-        ret_code = torch_converter.convert_weights_to_pytorch_script(model_rdf, output_path, use_tracing)
+        ret_code = torch_converter.convert_weights_to_torchscript(model_rdf, output_path, use_tracing)
         sys.exit(ret_code)
 
-    convert_torch_weights_to_torchscript.__doc__ = torch_converter.convert_weights_to_pytorch_script.__doc__
+    convert_torch_weights_to_torchscript.__doc__ = torch_converter.convert_weights_to_torchscript.__doc__
 
 
 if keras_converter is not None:

bioimageio/core/build_spec/build_model.py

@@ -113,10 +113,8 @@ def _get_weights(
             source=weight_source, sha256=weight_hash, opset_version=opset_version, **attachments
         )
 
-    elif weight_type == "pytorch_script":
-        weights = model_spec.raw_nodes.PytorchScriptWeightsEntry(
-            source=weight_source, sha256=weight_hash, **attachments
-        )
+    elif weight_type == "torchscript":
+        weights = model_spec.raw_nodes.TorchscriptWeightsEntry(source=weight_source, sha256=weight_hash, **attachments)
 
     elif weight_type == "keras_hdf5":
         if tensorflow_version is None:

bioimageio/core/prediction_pipeline/_model_adapters/_model_adapter.py

@@ -7,7 +7,7 @@
 
 #: Known weight formats in order of priority
 #: First match wins
-_WEIGHT_FORMATS = ["pytorch_state_dict", "tensorflow_saved_model_bundle", "pytorch_script", "onnx", "keras_hdf5"]
+_WEIGHT_FORMATS = ["pytorch_state_dict", "tensorflow_saved_model_bundle", "torchscript", "onnx", "keras_hdf5"]
 
 
 class ModelAdapter(abc.ABC):
@@ -114,7 +114,7 @@ def create_model_adapter(
             return adapter_cls(bioimageio_model=bioimageio_model, devices=devices)
 
     raise RuntimeError(
-        f"weight format {weight_format} not among weight formats listed in model: {list(bioimageio_model.weights.keys())}"
+        f"weight format {weight_format} not among formats listed in model: {list(bioimageio_model.weights.keys())}"
     )
 
 
@@ -139,7 +139,7 @@ def _get_model_adapter(weight_format: str) -> Type[ModelAdapter]:
 
         return ONNXModelAdapter
 
-    elif weight_format == "pytorch_script":
+    elif weight_format == "torchscript":
         from ._torchscript_model_adapter import TorchscriptModelAdapter
 
         return TorchscriptModelAdapter

bioimageio/core/prediction_pipeline/_model_adapters/_torchscript_model_adapter.py

@@ -11,7 +11,7 @@
 
 class TorchscriptModelAdapter(ModelAdapter):
     def _load(self, *, devices: Optional[List[str]] = None):
-        weight_path = str(self.bioimageio_model.weights["pytorch_script"].source.resolve())
+        weight_path = str(self.bioimageio_model.weights["torchscript"].source.resolve())
         if devices is None:
             self.devices = ["cuda" if torch.cuda.is_available() else "cpu"]
         else:

bioimageio/core/resource_io/nodes.py

@@ -161,7 +161,7 @@ class PytorchStateDictWeightsEntry(Node, model_raw_nodes.PytorchStateDictWeights
 
 
 @dataclass
-class PytorchScriptWeightsEntry(Node, model_raw_nodes.PytorchScriptWeightsEntry):
+class TorchscriptWeightsEntry(Node, model_raw_nodes.TorchscriptWeightsEntry):
     source: Path = missing
 
 
@@ -184,10 +184,10 @@ class Attachments(Node, model_raw_nodes.Attachments):
 WeightsEntry = Union[
     KerasHdf5WeightsEntry,
     OnnxWeightsEntry,
-    PytorchScriptWeightsEntry,
     PytorchStateDictWeightsEntry,
     TensorflowJsWeightsEntry,
     TensorflowSavedModelBundleWeightsEntry,
+    TorchscriptWeightsEntry,
 ]
 
 

bioimageio/core/weight_converter/torch/__init__.py

@@ -1,2 +1,2 @@
 from .onnx import convert_weights_to_onnx
-from .torchscript import convert_weights_to_pytorch_script
+from .torchscript import convert_weights_to_torchscript

bioimageio/core/weight_converter/torch/torchscript.py

@@ -72,7 +72,7 @@ def _check(input_):
             return ret
 
 
-def convert_weights_to_pytorch_script(
+def convert_weights_to_torchscript(
     model_spec: Union[str, Path, spec.model.raw_nodes.Model], output_path: Union[str, Path], use_tracing: bool = True
 ):
     """Convert model weights from format 'pytorch_state_dict' to 'torchscript'.

tests/build_spec/test_build_spec.py

@@ -28,10 +28,10 @@ def _test_build_spec(
         model_kwargs = None if weight_spec.kwargs is missing else weight_spec.kwargs
         architecture = str(weight_spec.architecture)
         weight_type_ = None  # the weight type can be auto-detected
-    elif weight_type == "pytorch_script":
+    elif weight_type == "torchscript":
         architecture = None
         model_kwargs = None
-        weight_type_ = "pytorch_script"  # the weight type CANNOT be auto-detcted
+        weight_type_ = "torchscript"  # the weight type CANNOT be auto-detcted
     else:
         architecture = None
         model_kwargs = None
@@ -109,7 +109,7 @@ def test_build_spec_implicit_output_shape(unet2d_nuclei_broad_model, tmp_path):
 
 
 def test_build_spec_torchscript(any_torchscript_model, tmp_path):
-    _test_build_spec(any_torchscript_model, tmp_path / "model.zip", "pytorch_script")
+    _test_build_spec(any_torchscript_model, tmp_path / "model.zip", "torchscript")
 
 
 def test_build_spec_onnx(any_onnx_model, tmp_path):
@@ -133,7 +133,7 @@ def test_build_spec_tfjs(any_tensorflow_js_model, tmp_path):
 
 
 def test_build_spec_deepimagej(unet2d_nuclei_broad_model, tmp_path):
-    _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "pytorch_script", add_deepimagej_config=True)
+    _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "torchscript", add_deepimagej_config=True)
 
 
 def test_build_spec_deepimagej_keras(unet2d_keras, tmp_path):

tests/prediction_pipeline/test_device_management.py

@@ -1,5 +1,3 @@
-from functools import wraps
-
 import numpy as np
 import pytest
 import xarray as xr
@@ -60,7 +58,7 @@ def test_device_management_torch(any_torch_model):
 
 @skip_on(TooFewDevicesException, reason="Too few devices")
 def test_device_management_torchscript(any_torchscript_model):
-    _test_device_management(any_torchscript_model, "pytorch_script")
+    _test_device_management(any_torchscript_model, "torchscript")
 
 
 @pytest.mark.skipif(pytest.skip_torch, reason="requires torch for device discovery")