2542 add use_deterministic_algorithms flag

monai/utils/__init__.py

@@ -48,6 +48,7 @@
     fall_back_tuple,
     first,
     get_seed,
+    has_option,
     is_scalar,
     is_scalar_tensor,
     issequenceiterable,
@@ -66,7 +67,6 @@
     get_full_type_name,
     get_package_version,
     get_torch_version_tuple,
-    has_option,
     load_submodules,
     min_version,
     optional_import,

monai/utils/misc.py

@@ -22,6 +22,8 @@
 import numpy as np
 import torch
 
+from monai.utils.module import get_torch_version_tuple
+
 __all__ = [
     "zip_with",
     "star_zip_with",
@@ -214,6 +216,7 @@ def get_seed() -> Optional[int]:
 
 def set_determinism(
     seed: Optional[int] = np.iinfo(np.uint32).max,
+    use_deterministic_algorithms: Optional[bool] = None,
     additional_settings: Optional[Union[Sequence[Callable[[int], Any]], Callable[[int], Any]]] = None,
 ) -> None:
     """
@@ -224,8 +227,8 @@ def set_determinism(
             It is recommended to set a large seed, i.e. a number that has a good balance
             of 0 and 1 bits. Avoid having many 0 bits in the seed.
             if set to None, will disable deterministic training.
-        additional_settings: additional settings
-            that need to set random seed.
+        use_deterministic_algorithms: Set whether PyTorch operations must use "deterministic" algorithms.
+        additional_settings: additional settings that need to set random seed.
 
     """
     if seed is None:
@@ -254,6 +257,15 @@ def set_determinism(
         torch.backends.cudnn.deterministic = _flag_deterministic
         torch.backends.cudnn.benchmark = _flag_cudnn_benchmark
 
+    if use_deterministic_algorithms is not None:
+        torch_ver = get_torch_version_tuple()
+        if torch_ver >= (1, 9):
+            torch.use_deterministic_algorithms(use_deterministic_algorithms)
+        elif torch_ver >= (1, 7):
+            torch.set_deterministic(use_deterministic_algorithms)  # beta feature
+        else:
+            warnings.warn("use_deterministic_algorithms=True, but PyTorch version is too old to set the mode.")
+
 
 def list_to_dict(items):
     """
@@ -359,3 +371,13 @@ class ImageMetaKey:
 
     FILENAME_OR_OBJ = "filename_or_obj"
     PATCH_INDEX = "patch_index"
+
+
+def has_option(obj, keywords: Union[str, Sequence[str]]) -> bool:
+    """
+    Return a boolean indicating whether the given callable `obj` has the `keywords` in its signature.
+    """
+    if not callable(obj):
+        return False
+    sig = inspect.signature(obj)
+    return all(key in sig.parameters for key in ensure_tuple(keywords))

monai/utils/module.py

@@ -9,18 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 import sys
 import warnings
 from importlib import import_module
 from pkgutil import walk_packages
 from re import match
-from typing import Any, Callable, List, Sequence, Tuple, Union
+from typing import Any, Callable, List, Tuple
 
 import torch
 
-from .misc import ensure_tuple
-
 OPTIONAL_IMPORT_MSG_FMT = "{}"
 
 __all__ = [
@@ -32,7 +29,6 @@
     "optional_import",
     "load_submodules",
     "get_full_type_name",
-    "has_option",
     "get_package_version",
     "get_torch_version_tuple",
     "PT_BEFORE_1_7",
@@ -243,16 +239,6 @@ def __call__(self, *_args, **_kwargs):
     return _LazyRaise(), False
 
 
-def has_option(obj, keywords: Union[str, Sequence[str]]) -> bool:
-    """
-    Return a boolean indicating whether the given callable `obj` has the `keywords` in its signature.
-    """
-    if not callable(obj):
-        return False
-    sig = inspect.signature(obj)
-    return all(key in sig.parameters for key in ensure_tuple(keywords))
-
-
 def get_package_version(dep_name, default="NOT INSTALLED or UNKNOWN VERSION."):
     """
     Try to load package and get version. If not found, return `default`.

tests/test_set_determinism.py

@@ -15,6 +15,7 @@
 import torch
 
 from monai.utils import get_seed, set_determinism
+from tests.utils import SkipIfBeforePyTorchVersion, skip_if_no_cuda
 
 
 class TestSetDeterminism(unittest.TestCase):
@@ -49,5 +50,21 @@ def test_values(self):
         set_determinism(seed=None)
 
 
+class TestSetFlag(unittest.TestCase):
+    def setUp(self):
+        set_determinism(1, use_deterministic_algorithms=True)
+
+    @SkipIfBeforePyTorchVersion((1, 8))  # beta feature
+    @skip_if_no_cuda
+    def test_algo(self):
+        with self.assertRaises(RuntimeError):
+            x = torch.randn(20, 16, 50, 44, 31, requires_grad=True, device="cuda:0")
+            y = torch.nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))(x)
+            y.sum().backward()
+
+    def tearDown(self):
+        set_determinism(None, use_deterministic_algorithms=False)
+
+
 if __name__ == "__main__":
     unittest.main()