"""Seed utilities for reproducibility.
Authors
* Adel Moumen 2024
"""
import os
import random
import torch
from speechbrain.utils.distributed import get_rank, rank_prefixed_message
from speechbrain.utils.logger import get_logger
logger = get_logger(__name__)
max_seed_value = 4294967295 # 2^32 - 1 (uint32)
min_seed_value = 0
[docs]
def seed_everything(
seed: int = 0, verbose: bool = True, deterministic: bool = False
) -> int:
r"""Function that sets the seed for pseudo-random number generators in: torch, numpy, and Python's random module.
Arguments
---------
seed: int
the integer value seed for global random state.
verbose: bool
Whether to print a message on each rank with the seed being set.
deterministic: bool
Whether to set the seed for deterministic operations.
Returns
-------
int
The seed that was set.
"""
# if DDP, we need to offset the seed by the rank
# to avoid having the same seed on all processes
seed_offset = 0 if get_rank() is None else get_rank()
if not (min_seed_value <= seed <= max_seed_value):
logger.info(
f"{seed} is not in bounds, numpy accepts from {min_seed_value} to {max_seed_value}",
)
seed = seed_offset
else:
seed += seed_offset
if verbose:
logger.info(
rank_prefixed_message(f"Setting seed to {seed}"),
main_process_only=False,
)
os.environ["SB_GLOBAL_SEED"] = str(seed)
random.seed(seed)
# if numpy is available, seed it
try:
import numpy as np
np.random.seed(seed)
except ImportError:
pass
torch.manual_seed(seed)
# safe to call this function even if cuda is not available
torch.cuda.manual_seed_all(seed)
if deterministic:
torch.use_deterministic_algorithms(True)
return seed