Skip to content

edspdf.utils.random

set_seed

Source code in edspdf/utils/random.py
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class set_seed:
    def __init__(self, seed, cuda=torch.cuda.is_available()):
        """
        Set seed values for random generators.
        If used as a context, restore the random state
        used before entering the context.

        Parameters
        ----------
        seed: int
            Value used as a seed.
        cuda: bool
            Saves the cuda random states too
        """
        # if seed is True:
        #     seed = random.randint(1, 2**16)
        if seed is True:
            seed = random.randint(1, 2**16)
        self.state = get_random_generator_state(cuda)
        if seed is not None:
            random.seed(seed)
            torch.manual_seed(seed)
            np.random.seed(seed)
            if cuda:  # pragma: no cover
                torch.cuda.manual_seed(seed)
                torch.cuda.manual_seed_all(seed)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        set_random_generator_state(self.state)

__init__(seed, cuda=torch.cuda.is_available())

Set seed values for random generators. If used as a context, restore the random state used before entering the context.

PARAMETER DESCRIPTION
seed

Value used as a seed.

cuda

Saves the cuda random states too

DEFAULT: torch.cuda.is_available()

Source code in edspdf/utils/random.py
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(self, seed, cuda=torch.cuda.is_available()):
    """
    Set seed values for random generators.
    If used as a context, restore the random state
    used before entering the context.

    Parameters
    ----------
    seed: int
        Value used as a seed.
    cuda: bool
        Saves the cuda random states too
    """
    # if seed is True:
    #     seed = random.randint(1, 2**16)
    if seed is True:
        seed = random.randint(1, 2**16)
    self.state = get_random_generator_state(cuda)
    if seed is not None:
        random.seed(seed)
        torch.manual_seed(seed)
        np.random.seed(seed)
        if cuda:  # pragma: no cover
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)

get_random_generator_state(cuda=torch.cuda.is_available())

Get the torch, numpy and random random generator state.

PARAMETER DESCRIPTION
cuda

Saves the cuda random states too

DEFAULT: torch.cuda.is_available()

RETURNS DESCRIPTION
RandomGeneratorState
Source code in edspdf/utils/random.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def get_random_generator_state(cuda=torch.cuda.is_available()):
    """
    Get the `torch`, `numpy` and `random` random generator state.
    Parameters
    ----------
    cuda: bool
        Saves the cuda random states too

    Returns
    -------
    RandomGeneratorState
    """
    return RandomGeneratorState(
        random.getstate(),
        torch.random.get_rng_state(),
        np.random.get_state(),
        torch.cuda.get_rng_state_all() if cuda else None,
    )

set_random_generator_state(state)

Set the torch, numpy and random random generator state.

PARAMETER DESCRIPTION
state

Source code in edspdf/utils/random.py
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def set_random_generator_state(state):
    """
    Set the `torch`, `numpy` and `random` random generator state.
    Parameters
    ----------
    state: RandomGeneratorState
    """
    random.setstate(state.random)
    torch.random.set_rng_state(state.torch)
    np.random.set_state(state.numpy)
    if (
        state.torch_cuda is not None
        and torch.cuda.is_available()
        and len(state.torch_cuda) == torch.cuda.device_count()
    ):  # pragma: no cover
        torch.cuda.set_rng_state_all(state.torch_cuda)