跳到内容

随机实用程序

get_seed_from_rng(rng, dtype=np.int64)

从现有的随机生成器生成确定性的随机种子。

这特别有用,因为设置 torch 种子时,它不接受数字元组,但我们在使用 epoch、索引和全局种子初始化 numpy 随机生成器时经常这样做。

用于从 numpy 随机生成器为 torch 随机生成器播种。

源代码位于 bionemo/core/utils/random_utils.py
52
53
54
55
56
57
58
59
60
def get_seed_from_rng(rng: np.random.Generator, dtype: Type[np.signedinteger] = np.int64) -> int:
    """Generates a deterministic random seed from an existing random generator.

    This is useful in particular because setting the torch seed doesn't want to accept a tuple of numbers, we we often
    do in initializing a numpy random generator with epoch, index, and global seeds.

    Used to seed a torch random generator from a numpy random generator.
    """
    return int(rng.integers(np.iinfo(dtype).max))

random_numpy_context(seed=42)

用于设置 numpy 随机状态的上下文管理器。

状态在进入时保存,并在退出时恢复到之前的状态。这样,你可以在使用此函数的 with 上下文中运行需要随机状态的代码,并恢复到之前的任何状态。这对于测试非常有用,因为您不希望一个测试的随机状态影响其他测试。

示例

import numpy as np from bionemo.core.utils.random_utils import random_numpy_context ori_state = np.random.get_state() with random_numpy_context(45): np.random.randint(5) # this will change the state new_state = np.random.get_state() assert ori_state == new_state

源代码位于 bionemo/core/utils/random_utils.py
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@contextmanager
def random_numpy_context(seed: int = 42) -> Iterator[None]:
    """Context manager for setting numpy random state.

    The state is saved on entry and restored on exit to what it was. This way you can run code that needs random state
    in a `with` context using this function, and get back to whatever state was there before. This is useful for testing
    where you don't want the random state from one test to impact other tests.

    Example:
        >>> import numpy as np
        >>> from bionemo.core.utils.random_utils import random_numpy_context
        >>> ori_state = np.random.get_state()
        >>> with random_numpy_context(45):
            np.random.randint(5) # this will change the state
        >>> new_state = np.random.get_state()
        >>> assert ori_state == new_state
    """
    state = np.random.get_state()  # just fail if this fails
    try:
        np.random.seed(seed)
        yield
    finally:
        np.random.set_state(state)