add enable_determinism

Author: FynnBe

bioimageio/core/__init__.py

@@ -33,7 +33,12 @@
     tensor,
 )
 from ._prediction_pipeline import PredictionPipeline, create_prediction_pipeline
-from ._resource_tests import load_description_and_test, test_description, test_model
+from ._resource_tests import (
+    enable_determinism,
+    load_description_and_test,
+    test_description,
+    test_model,
+)
 from ._settings import settings
 from .axis import Axis, AxisId
 from .block_meta import BlockMeta
@@ -71,6 +76,7 @@
     "create_prediction_pipeline",
     "digest_spec",
     "dump_description",
+    "enable_determinism",
     "io",
     "load_dataset_description",
     "load_description_and_test",

bioimageio/core/_resource_tests.py

@@ -32,6 +32,67 @@
 from .utils import VERSION
 
 
+def enable_determinism(mode: Literal["seed_only", "full"]):
+    """Seed and configure ML frameworks for maximum reproducibility.
+    May degrade performance. Only recommended for testing reproducibility!
+
+    Seed any random generators and (if **mode**=="full") request ML frameworks to use
+    deterministic algorithms.
+    Notes:
+        - **mode** == "full"  might degrade performance and throw exceptions.
+        - Subsequent inference calls might still differ. Call before each function
+          (sequence) that is expected to be reproducible.
+        - Degraded performance: Use for testing reproducibility only!
+        - Recipes:
+            - [PyTorch](https://pytorch.org/docs/stable/notes/randomness.html)
+            - [Keras](https://keras.io/examples/keras_recipes/reproducibility_recipes/)
+            - [NumPy](https://numpy.org/doc/2.0/reference/random/generated/numpy.random.seed.html)
+    """
+    try:
+        try:
+            import numpy.random
+        except ImportError:
+            pass
+        else:
+            numpy.random.seed(0)
+    except Exception as e:
+        logger.debug(str(e))
+
+    try:
+        try:
+            import torch
+        except ImportError:
+            pass
+        else:
+            _ = torch.manual_seed(0)
+            torch.use_deterministic_algorithms(mode == "full")
+    except Exception as e:
+        logger.debug(str(e))
+
+    try:
+        try:
+            import keras
+        except ImportError:
+            pass
+        else:
+            keras.utils.set_random_seed(0)
+    except Exception as e:
+        logger.debug(str(e))
+
+    try:
+        try:
+            import tensorflow as tf  # pyright: ignore[reportMissingImports]
+        except ImportError:
+            pass
+        else:
+            tf.random.seed(0)
+            if mode == "full":
+                tf.config.experimental.enable_op_determinism()
+            # TODO: find possibility to switch it off again??
+    except Exception as e:
+        logger.debug(str(e))
+
+
 def test_model(
     source: Union[v0_5.ModelDescr, PermissiveFileSource],
     weight_format: Optional[WeightsFormat] = None,

tests/test_resource_tests.py

@@ -1,6 +1,17 @@
+from typing import Literal
+
+import pytest
+
 from bioimageio.spec import InvalidDescr
 
 
+@pytest.mark.parametrize("mode", ["seed_only", "full"])
+def test_enable_determinism(mode: Literal["seed_only", "full"]):
+    from bioimageio.core import enable_determinism
+
+    enable_determinism(mode)
+
+
 def test_error_for_wrong_shape(stardist_wrong_shape: str):
     from bioimageio.core._resource_tests import test_model