Module `xingyun.random.set_random_seed`

Expand source code

import random
import time
from typing import Literal, TypeAlias, NoReturn, List
from xingyun.universal.import_module import my_import_module

'''literals that are allowed in the `module` parameter.'''
RandomAllowedModule: TypeAlias = Literal["torch"] | Literal["numpy"] | Literal["random"] 

def set_module_seed(seed: int , module: RandomAllowedModule):
    '''set random seed of a specific module.'''
    
    if module == "torch":
        torch = my_import_module("torch")
        cuda  = my_import_module("torch.cuda")

        if torch is None or cuda is None:
            return 

        torch.manual_seed(seed)
        cuda.manual_seed_all(seed)

        backends = my_import_module("torch.backends")
        backends.cudnn.deterministic = True
        backends.cudnn.benchmark = False
    if module == "numpy":
        np = my_import_module("numpy")

        if np is None:
            return 
        
        np.random.seed(seed)
        
    if module == "random":
        random.seed(seed)


def set_random_seed(seed: int , modules: List[RandomAllowedModule] = ["random" , "torch" , "numpy"]):
    '''Set random seed.
        
    ### Parameters
        - seed: random seed.
        - modules: Which module to set. Support `random`, `torch` and `numpy`.
            Notice that the random seed of `torch.cuda` will also be set.
    '''
    [set_module_seed(seed, m) for m in modules]

Functions

def set_module_seed(seed: int, module: Union[Literal['torch'], Literal['numpy'], Literal['random']])

set random seed of a specific module.

Expand source code

def set_module_seed(seed: int , module: RandomAllowedModule):
    '''set random seed of a specific module.'''
    
    if module == "torch":
        torch = my_import_module("torch")
        cuda  = my_import_module("torch.cuda")

        if torch is None or cuda is None:
            return 

        torch.manual_seed(seed)
        cuda.manual_seed_all(seed)

        backends = my_import_module("torch.backends")
        backends.cudnn.deterministic = True
        backends.cudnn.benchmark = False
    if module == "numpy":
        np = my_import_module("numpy")

        if np is None:
            return 
        
        np.random.seed(seed)
        
    if module == "random":
        random.seed(seed)

def set_random_seed(seed: int, modules: List[Union[Literal['torch'], Literal['numpy'], Literal['random']]] = ['random', 'torch', 'numpy'])

Set random seed.

Parameters

- seed: random seed.
- modules: Which module to set. Support <code>random</code>, <code>torch</code> and <code>numpy</code>.
    Notice that the random seed of <code>torch.cuda</code> will also be set.

Expand source code

def set_random_seed(seed: int , modules: List[RandomAllowedModule] = ["random" , "torch" , "numpy"]):
    '''Set random seed.
        
    ### Parameters
        - seed: random seed.
        - modules: Which module to set. Support `random`, `torch` and `numpy`.
            Notice that the random seed of `torch.cuda` will also be set.
    '''
    [set_module_seed(seed, m) for m in modules]