Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,380 Bytes
137645c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import os
import random
import numpy as np
import torch
def seed_everything(seed=0, deterministic_ops=True, allow_tf32=False):
"""
Sets the seed for reproducibility across various libraries and frameworks, and configures PyTorch backend settings.
Args:
seed (int): The seed value for random number generation. Default is 0.
deterministic_ops (bool): Whether to enable deterministic operations in PyTorch.
Enabling this can make results reproducible at the cost of potential performance degradation. Default is True.
allow_tf32 (bool): Whether to allow TensorFloat-32 (TF32) precision in PyTorch operations. TF32 can improve performance but may affect reproducibility. Default is False.
Effects:
- Seeds Python's random module, NumPy, and PyTorch (CPU and GPU).
- Sets the environment variable `PYTHONHASHSEED` to the specified seed.
- Configures PyTorch to use deterministic algorithms if `deterministic_ops` is True.
- Configures TensorFloat-32 precision based on `allow_tf32`.
- Issues warnings if configurations may impact reproducibility.
Notes:
- Setting `torch.backends.cudnn.deterministic` to False allows nondeterministic operations, which may introduce variability.
- Allowing TF32 (`allow_tf32=True`) may lead to non-reproducible results, especially in matrix operations.
"""
# Seed standard random number generators
random.seed(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
# Seed PyTorch random number generators
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# Configure deterministic operations
if deterministic_ops:
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True)
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
else:
torch.backends.cudnn.deterministic = False
print("WARNING: torch.backends.cudnn.deterministic is set to False, reproducibility is not guaranteed.")
# Configure TensorFloat-32 precision
if allow_tf32:
print("WARNING: TensorFloat-32 (TF32) is enabled; reproducibility is not guaranteed.")
torch.backends.cudnn.allow_tf32 = allow_tf32 # Default True in PyTorch 2.6.0
torch.backends.cuda.matmul.allow_tf32 = allow_tf32 # Default False in PyTorch 2.6.0
|