Merge pull request #176 from bioimage-io/more-tf-features

More tf features

Author: constantinpape

.github/workflows/build.yml

@@ -80,6 +80,7 @@ jobs:
       uses: conda-incubator/setup-miniconda@v2
       with:
         auto-update-conda: true
+        mamba-version: "*"
         channel-priority: strict
         activate-environment: bio-core-tf
         environment-file: dev/environment-tf.yaml

bioimageio/core/__main__.py

@@ -23,6 +23,11 @@
 except ImportError:
     torch_converter = None
 
+try:
+    from bioimageio.core.weight_converter import keras as keras_converter
+except ImportError:
+    keras_converter = None
+
 
 # extend help/version string by core version
 help_version_core = f"bioimageio.core {__version__}"
@@ -231,7 +236,8 @@ def convert_torch_weights_to_onnx(
         use_tracing: bool = typer.Option(True, help="Whether to use torch.jit tracing or scripting."),
         verbose: bool = typer.Option(True, help="Verbosity"),
     ) -> int:
-        return torch_converter.convert_weights_to_onnx(model_rdf, output_path, opset_version, use_tracing, verbose)
+        ret_code = torch_converter.convert_weights_to_onnx(model_rdf, output_path, opset_version, use_tracing, verbose)
+        sys.exit(ret_code)
 
     convert_torch_weights_to_onnx.__doc__ = torch_converter.convert_weights_to_onnx.__doc__
 
@@ -249,5 +255,22 @@ def convert_torch_weights_to_torchscript(
     convert_torch_weights_to_torchscript.__doc__ = torch_converter.convert_weights_to_pytorch_script.__doc__
 
 
+if keras_converter is not None:
+
+    @app.command()
+    def convert_keras_weights_to_tensorflow(
+        model_rdf: Path = typer.Argument(
+            ..., help="Path to the model resource description file (rdf.yaml) or zipped model."
+        ),
+        output_path: Path = typer.Argument(..., help="Where to save the tensorflow weights."),
+    ) -> int:
+        ret_code = keras_converter.convert_weights_to_tensorflow_saved_model_bundle(model_rdf, output_path)
+        sys.exit(ret_code)
+
+    convert_keras_weights_to_tensorflow.__doc__ = (
+        keras_converter.convert_weights_to_tensorflow_saved_model_bundle.__doc__
+    )
+
+
 if __name__ == "__main__":
     app()

bioimageio/core/build_spec/build_model.py

@@ -78,7 +78,16 @@ def _get_pytorch_state_dict_weight_kwargs(architecture, model_kwargs, root):
     return weight_kwargs, tmp_archtecture
 
 
-def _get_weights(original_weight_source, weight_type, root, architecture=None, model_kwargs=None, **kwargs):
+def _get_weights(
+    original_weight_source,
+    weight_type,
+    root,
+    architecture=None,
+    model_kwargs=None,
+    tensorflow_version=None,
+    opset_version=None,
+    **kwargs,
+):
     weight_path = resolve_source(original_weight_source, root)
     if weight_type is None:
         weight_type = _infer_weight_type(weight_path)
@@ -98,8 +107,10 @@ def _get_weights(original_weight_source, weight_type, root, architecture=None, m
         )
 
     elif weight_type == "onnx":
+        if opset_version is None:
+            raise ValueError("opset_version needs to be passed for building an onnx model")
         weights = model_spec.raw_nodes.OnnxWeightsEntry(
-            source=weight_source, sha256=weight_hash, opset_version=kwargs.get("opset_version", 12), **attachments
+            source=weight_source, sha256=weight_hash, opset_version=opset_version, **attachments
         )
 
     elif weight_type == "pytorch_script":
@@ -108,26 +119,32 @@ def _get_weights(original_weight_source, weight_type, root, architecture=None, m
         )
 
     elif weight_type == "keras_hdf5":
+        if tensorflow_version is None:
+            raise ValueError("tensorflow_version needs to be passed for building a keras model")
         weights = model_spec.raw_nodes.KerasHdf5WeightsEntry(
             source=weight_source,
             sha256=weight_hash,
-            tensorflow_version=kwargs.get("tensorflow_version", "1.15"),
+            tensorflow_version=tensorflow_version,
             **attachments,
         )
 
     elif weight_type == "tensorflow_saved_model_bundle":
+        if tensorflow_version is None:
+            raise ValueError("tensorflow_version needs to be passed for building a tensorflow model")
         weights = model_spec.raw_nodes.TensorflowSavedModelBundleWeightsEntry(
             source=weight_source,
             sha256=weight_hash,
-            tensorflow_version=kwargs.get("tensorflow_version", "1.15"),
+            tensorflow_version=tensorflow_version,
             **attachments,
         )
 
     elif weight_type == "tensorflow_js":
+        if tensorflow_version is None:
+            raise ValueError("tensorflow_version needs to be passed for building a tensorflow_js model")
         weights = model_spec.raw_nodes.TensorflowJsWeightsEntry(
             source=weight_source,
             sha256=weight_hash,
-            tensorflow_version=kwargs.get("tensorflow_version", "1.15"),
+            tensorflow_version=tensorflow_version,
             **attachments,
         )
 
@@ -471,6 +488,8 @@ def build_model(
     links: Optional[List[str]] = None,
     root: Optional[Union[Path, str]] = None,
     add_deepimagej_config: bool = False,
+    tensorflow_version: Optional[str] = None,
+    opset_version: Optional[int] = None,
     **weight_kwargs,
 ):
     """Create a zipped bioimage.io model.
@@ -539,7 +558,10 @@ def build_model(
         dependencies: relative path to file with dependencies for this model.
         root: optional root path for relative paths. This can be helpful when building a spec from another model spec.
         add_deepimagej_config: add the deepimagej config to the model.
-        weight_kwargs: keyword arguments for this weight type, e.g. "tensorflow_version".
+        tensorflow_version: the tensorflow version used for training the model.
+            Needs to be passed for tensorflow or keras models.
+        opset_version: the opset version used in this model. Needs to be passed for onnx models.
+        weight_kwargs: additional keyword arguments for this weight type.
     """
     if root is None:
         root = "."
@@ -624,7 +646,16 @@ def build_model(
     covers = _ensure_local(covers, root)
 
     # parse the weights
-    weights, tmp_archtecture = _get_weights(weight_uri, weight_type, root, architecture, model_kwargs, **weight_kwargs)
+    weights, tmp_archtecture = _get_weights(
+        weight_uri,
+        weight_type,
+        root,
+        architecture,
+        model_kwargs,
+        tensorflow_version=tensorflow_version,
+        opset_version=opset_version,
+        **weight_kwargs,
+    )
 
     # validate the sample inputs and outputs (if given)
     if sample_inputs is not None:
@@ -732,11 +763,24 @@ def add_weights(
     weight_uri: Union[str, Path],
     weight_type: Optional[str] = None,
     output_path: Optional[Union[str, Path]] = None,
+    architecture: Optional[str] = None,
+    model_kwargs: Optional[Dict[str, Union[int, float, str]]] = None,
+    tensorflow_version: Optional[str] = None,
+    opset_version: Optional[str] = None,
     **weight_kwargs,
 ):
     """Add weight entry to bioimage.io model."""
     # we need to pass the weight path as abs path to avoid confusion with different root directories
-    new_weights, tmp_arch = _get_weights(Path(weight_uri).absolute(), weight_type, root=Path("."), **weight_kwargs)
+    new_weights, tmp_arch = _get_weights(
+        Path(weight_uri).absolute(),
+        weight_type,
+        root=Path("."),
+        architecture=architecture,
+        model_kwargs=model_kwargs,
+        tensorflow_version=tensorflow_version,
+        opset_version=opset_version,
+        **weight_kwargs,
+    )
     model.weights.update(new_weights)
     if output_path is not None:
         model_package = export_resource_package(model, output_path=output_path)

bioimageio/core/prediction_pipeline/_model_adapters/_keras_model_adapter.py

@@ -4,15 +4,33 @@
 # by default, we use the keras integrated with tensorflow
 try:
     from tensorflow import keras
+    import tensorflow as tf
+
+    TF_VERSION = tf.__version__
 except Exception:
     import keras
+
+    TF_VERSION = None
 import xarray as xr
 
 from ._model_adapter import ModelAdapter
 
 
 class KerasModelAdapter(ModelAdapter):
     def _load(self, *, devices: Optional[Sequence[str]] = None) -> None:
+        try:
+            model_tf_version = self.bioimageio_model.weights[self.weight_format].tensorflow_version.version
+        except AttributeError:
+            model_tf_version = None
+
+        if TF_VERSION is None or model_tf_version is None:
+            warnings.warn("Could not check tensorflow versions. The prediction results may be wrong.")
+        elif tuple(model_tf_version[:2]) != tuple(map(int, TF_VERSION.split(".")))[:2]:
+            warnings.warn(
+                f"Model tensorflow version {model_tf_version} does not match {TF_VERSION}."
+                "The prediction results may be wrong"
+            )
+
         # TODO keras device management
         if devices is not None:
             warnings.warn(f"Device management is not implemented for keras yet, ignoring the devices {devices}")

bioimageio/core/prediction_pipeline/_model_adapters/_tensorflow_model_adapter.py

@@ -36,10 +36,24 @@ def _load_model(self, weight_file):
 
     def _load(self, *, devices: Optional[List[str]] = None):
         try:
-            tf_version = self.bioimageio_model.weights[self.weight_format].tensorflow_version.version
+            model_tf_version = self.bioimageio_model.weights[self.weight_format].tensorflow_version.version
         except AttributeError:
-            tf_version = (1, 14, 0)
-        tf_major_ver = tf_version[0]
+            model_tf_version = None
+
+        tf_version = tf.__version__
+        tf_major_and_minor = tuple(map(int, tf_version.split(".")))[:2]
+        if model_tf_version is None:
+            warnings.warn(
+                "The model did not contain metadata about the tensorflow version used for training."
+                f"Cannot check if it is compatible with tf {tf_version}. The prediction result may be wrong."
+            )
+        elif tuple(model_tf_version[:2]) != tf_major_and_minor:
+            warnings.warn(
+                f"Model tensorflow version {model_tf_version} does not match {tf_version}."
+                "The prediction results may be wrong"
+            )
+
+        tf_major_ver = tf_major_and_minor[0]
         assert tf_major_ver in (1, 2)
         self.use_keras_api = tf_major_ver > 1 or self.weight_format == KerasModelAdapter.weight_format
 

bioimageio/core/weight_converter/keras/__init__.py

@@ -0,0 +1 @@
+from .tensorflow import convert_weights_to_tensorflow_saved_model_bundle

bioimageio/core/weight_converter/keras/tensorflow.py

@@ -0,0 +1,103 @@
+import os
+import shutil
+from pathlib import Path
+from typing import Union
+from zipfile import ZipFile
+
+import bioimageio.spec as spec
+from bioimageio.core import load_resource_description
+
+import tensorflow
+from tensorflow import saved_model
+
+
+# adapted from
+# https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236
+def _convert_tf1(keras_weight_path, output_path, zip_weights):
+    def build_tf_model():
+        keras_model = keras.models.load_model(keras_weight_path)
+
+        builder = saved_model.builder.SavedModelBuilder(output_path)
+        signature = saved_model.signature_def_utils.predict_signature_def(
+            inputs={"input": keras_model.input}, outputs={"output": keras_model.output}
+        )
+
+        signature_def_map = {saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature}
+
+        builder.add_meta_graph_and_variables(
+            keras.backend.get_session(), [saved_model.tag_constants.SERVING], signature_def_map=signature_def_map
+        )
+        builder.save()
+
+    try:
+        # try to build the tf model with the keras import from tensorflow
+        from tensorflow import keras
+        build_tf_model()
+    except Exception:
+        # if the above fails try to export with the standalone keras
+        import keras
+
+        build_tf_model()
+
+    if zip_weights:
+        zipped_model = f"{output_path}.zip"
+        # zip the weights
+        file_paths = []
+        for folder_names, subfolder, filenames in os.walk(os.path.join(output_path)):
+            for filename in filenames:
+                # create complete filepath of file in directory
+                file_paths.append(os.path.join(folder_names, filename))
+
+        with ZipFile(zipped_model, "w") as zip_obj:
+            for f in file_paths:
+                # Add file to zip
+                zip_obj.write(f, os.path.relpath(f, output_path))
+
+        try:
+            shutil.rmtree(output_path)
+        except Exception:
+            print("TensorFlow bundled model was not removed after compression")
+        print("TensorFlow model exported to", zipped_model)
+    else:
+        print("TensorFlow model exported to", output_path)
+    return 0
+
+
+def convert_weights_to_tensorflow_saved_model_bundle(
+    model_spec: Union[str, Path, spec.model.raw_nodes.Model], output_path: Union[str, Path]
+):
+    """Convert model weights from format 'keras_hdf5' to 'tensorflow_saved_model_bundle'.
+
+    Adapted from
+    https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py
+
+    Args:
+        model_spec: location of the resource for the input bioimageio model
+        output_path: where to save the tensorflow weights. This path must not exist yet.
+    """
+    tf_major_ver = int(tensorflow.__version__.split(".")[0])
+
+    path_ = Path(output_path)
+    if path_.suffix == ".zip":
+        path_ = Path(os.path.splitext(path_)[0])
+        zip_weights = True
+    else:
+        zip_weights = False
+
+    if path_.exists():
+        raise ValueError(f"The ouptut directory at {path_} must not exist.")
+
+    model = load_resource_description(model_spec)
+    assert "keras_hdf5" in model.weights
+    weight_spec = model.weights["keras_hdf5"]
+    weight_path = str(weight_spec.source)
+
+    if weight_spec.tensorflow_version:
+        model_tf_major_ver = weight_spec.tensorflow_version.version[0]
+        if model_tf_major_ver != tf_major_ver:
+            raise RuntimeError(f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}")
+
+    if tf_major_ver == 1:
+        return _convert_tf1(weight_path, str(path_), zip_weights)
+    else:
+        raise NotImplementedError

dev/environment-tf.yaml

@@ -12,7 +12,7 @@ dependencies:
   - pytest
   - python >=3.7,<3.8  # this environment is only available for python 3.7
   - xarray
-  - tensorflow >=1.12,<2.0
+  - tensorflow >1.14,<2.0
   - tifffile
   - pip:
     - keras==1.2.2