egrecho.utils.seeder.seed#

class egrecho.utils.seeder.seed.SeedWorkers(seed=42, rank=None, include_cuda=False)[source]#

Bases: object

Different workers with different seed.

If provide rank, it will randomlize acorssing node

Parameters:

  • seed (int) -- defaults to 42.

  • rank (int) -- defaults to None.

  • include_cuda (bool) -- whether fix randomlize cuda. defaults to False.

egrecho.utils.seeder.seed.isolate_rng(include_cuda=True)[source]#

A context manager that keeps track of the global random state, resets the global random state on exit to what it was before entering.

It supports isolating the states for PyTorch, Numpy, and Python built-in random number generators. referring to: https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/utilities/seed.py#isolate_rng

Parameters:

include_cuda (bool) -- Whether to allow this function to also control the torch.cuda random number generator. Set this to False when using the function in a forked process where CUDA re-initialization is prohibited.

Return type:

Generator[None, None, None]

Example

>>> import torch
>>> torch.manual_seed(1)  
<torch._C.Generator object at ...>
>>> with isolate_rng():
...     [torch.rand(1) for _ in range(3)]
[tensor([0.7576]), tensor([0.2793]), tensor([0.4031])]
>>> torch.rand(1)
tensor([0.7576])