diff --git a/src/careamics/utils/torch_utils.py b/src/careamics/utils/torch_utils.py index b2b9755fb..b46a253ff 100644 --- a/src/careamics/utils/torch_utils.py +++ b/src/careamics/utils/torch_utils.py @@ -4,7 +4,6 @@ These functions are used to control certain aspects and behaviours of PyTorch. """ import logging -import os import sys import torch @@ -48,46 +47,44 @@ def compile_model(model: torch.nn.Module) -> torch.nn.Module: return model -def seed_everything(seed: int) -> None: - """ - Seed all random number generators for reproducibility. - - Parameters - ---------- - seed : int - Seed. - """ - import random - - import numpy as np - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def setup_cudnn_reproducibility( - deterministic: bool = True, benchmark: bool = True -) -> None: - """ - Prepare CuDNN benchmark and sets it to be deterministic/non-deterministic mode. - - https://pytorch.org/docs/stable/notes/randomness.html#cuda-convolution-benchmarking. - - Parameters - ---------- - deterministic : bool - Deterministic mode, if running CuDNN backend. - benchmark : bool - If True, uses CuDNN heuristics to figure out which algorithm will be most - performant for your model architecture and input. False may slow down training. - """ - if torch.cuda.is_available(): - if deterministic: - deterministic = os.environ.get("CUDNN_DETERMINISTIC", "True") == "True" - torch.backends.cudnn.deterministic = deterministic - - if benchmark: - benchmark = os.environ.get("CUDNN_BENCHMARK", "True") == "True" - torch.backends.cudnn.benchmark = benchmark +# def seed_everything(seed: int) -> None: +# """ +# Seed all random number generators for reproducibility. + +# Parameters +# ---------- +# seed : int +# Seed. +# """ +# import random + +# import numpy as np + +# random.seed(seed) +# np.random.seed(seed) +# torch.manual_seed(seed) +# torch.cuda.manual_seed_all(seed) + + +# def setup_cudnn_reproducibility( +# deterministic: bool = True, benchmark: bool = True +# ) -> None: +# """ +# Prepare CuDNN benchmark and sets it to be deterministic/non-deterministic mode. + +# Parameters +# ---------- +# deterministic : bool +# Deterministic mode, if running CuDNN backend. +# benchmark : bool +# If True, uses CuDNN heuristics to figure out which algorithm will be most +# performant for your model architecture and input. False may slow down training +# """ +# if torch.cuda.is_available(): +# if deterministic: +# deterministic = os.environ.get("CUDNN_DETERMINISTIC", "True") == "True" +# torch.backends.cudnn.deterministic = deterministic + +# if benchmark: +# benchmark = os.environ.get("CUDNN_BENCHMARK", "True") == "True" +# torch.backends.cudnn.benchmark = benchmark diff --git a/tests/utils/test_torch_utils.py b/tests/utils/test_torch_utils.py index 8114eadb0..947a3685e 100644 --- a/tests/utils/test_torch_utils.py +++ b/tests/utils/test_torch_utils.py @@ -3,7 +3,6 @@ from careamics.utils.torch_utils import ( get_device, - setup_cudnn_reproducibility, ) @@ -14,10 +13,10 @@ def test_get_device(device): assert device.type == "cuda" if torch.cuda.is_available() else "cpu" -@pytest.mark.gpu -@pytest.mark.parametrize("deterministic", [True, False]) -@pytest.mark.parametrize("benchmark", [True, False]) -def test_setup_cudnn_reproducibility(deterministic, benchmark): - setup_cudnn_reproducibility(deterministic=deterministic, benchmark=benchmark) - assert torch.backends.cudnn.deterministic == deterministic - assert torch.backends.cudnn.benchmark == benchmark +# @pytest.mark.gpu +# @pytest.mark.parametrize("deterministic", [True, False]) +# @pytest.mark.parametrize("benchmark", [True, False]) +# def test_setup_cudnn_reproducibility(deterministic, benchmark): +# setup_cudnn_reproducibility(deterministic=deterministic, benchmark=benchmark) +# assert torch.backends.cudnn.deterministic == deterministic +# assert torch.backends.cudnn.benchmark == benchmark