Python torch.set_rng_state方法代码示例

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def seed_all_rng(seed=None):
    """
    Set the random seed for the RNG in torch, numpy and python.

    Args:
        seed (int): if None, will use a strong random seed.
    """
    if seed is None:
        seed = (
            os.getpid()
            + int(datetime.now().strftime("%S%f"))
            + int.from_bytes(os.urandom(2), "big")
        )
        logger = logging.getLogger(__name__)
        logger.info("Using a generated random seed {}".format(seed))
    np.random.seed(seed)
    torch.set_rng_state(torch.manual_seed(seed).get_state())
    random.seed(seed) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is class_index of the target class.
        """
        # create random image that is consistent with the index id
        rng_state = torch.get_rng_state()
        torch.manual_seed(index + self.random_offset)
        img = torch.randn(*self.image_size)
        target = torch.Tensor(1).random_(0, self.num_classes)[0]
        torch.set_rng_state(rng_state)

        # convert to PIL Image
        img = transforms.ToPILImage()(img)
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def restore_rng_states(device: torch.device,
                       rng_states: Deque[RNGStates],
                       ) -> Generator[None, None, None]:
    """:meth:`Recompute.backward` restores the random number generator states
    captured by :func:`save_rng_states` within its context.

    .. seealso:: :ref:`Referential Transparency`

    """
    cpu_rng_state, gpu_rng_state = rng_states.pop()

    gpu_devices: List[torch.device] = []
    if device.type == 'cuda':
        gpu_devices.append(device)

    with torch.random.fork_rng(gpu_devices):
        torch.set_rng_state(cpu_rng_state)
        if gpu_rng_state is not None:
            torch.cuda.set_rng_state(gpu_rng_state, device)
        yield 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def test_get_set_device_states(device, enabled):
    shape = (1, 1, 10, 10)
    if not torch.cuda.is_available() and device == 'cuda':
        pytest.skip('This test requires a GPU to be available')
    X = torch.ones(shape, device=device)
    devices, states = get_device_states(X)
    assert len(states) == (1 if device == 'cuda' else 0)
    assert len(devices) == (1 if device == 'cuda' else 0)
    cpu_rng_state = torch.get_rng_state()
    Y = X * torch.rand(shape, device=device)
    with torch.random.fork_rng(devices=devices, enabled=True):
        if enabled:
            if device == 'cpu':
                torch.set_rng_state(cpu_rng_state)
            else:
                set_device_states(devices=devices, states=states)
        Y2 = X * torch.rand(shape, device=device)
    assert torch.equal(Y, Y2) == enabled 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def backward(ctx, *args):
        if not torch.autograd._is_checkpoint_valid():
            raise RuntimeError("Checkpointing is not compatible with .grad(), please use .backward() if possible")
        inputs = ctx.saved_tensors
        # Stash the surrounding rng state, and mimic the state that was
        # present at this time during forward.  Restore the surrouding state
        # when we're done.
        rng_devices = [torch.cuda.current_device()] if ctx.had_cuda_in_fwd else []
        with torch.random.fork_rng(devices=rng_devices, enabled=preserve_rng_state):
            if preserve_rng_state:
                torch.set_rng_state(ctx.fwd_cpu_rng_state)
                if ctx.had_cuda_in_fwd:
                    torch.cuda.set_rng_state(ctx.fwd_cuda_rng_state)
            detached_inputs = detach_variable(inputs)
            with torch.enable_grad():
                outputs = ctx.run_function(*detached_inputs)

        if isinstance(outputs, torch.Tensor):
            outputs = (outputs,)
        torch.autograd.backward(outputs, args)
        return (None,) + tuple(inp.grad for inp in detached_inputs) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def load_states(self, epoch, device):
        """Loads the states from the checkpoint associated with ``epoch``.

        Args:
            epoch (int): ``--epoch`` command argument.
            device (str): ``--device`` command argument.
        """
        checkpoint_data = self.experiment.checkpoint_load(epoch, device)
        if isinstance(self.model, torch.nn.DataParallel):
            self.model.module.load_state_dict(checkpoint_data['model'])
        else:
            self.model.load_state_dict(checkpoint_data['model'])
        self.optimizer.load_state_dict(checkpoint_data['optimizer'])
        random.setstate(checkpoint_data['random_states'][0])
        np.random.set_state(checkpoint_data['random_states'][1])
        torch.set_rng_state(checkpoint_data['random_states'][2].cpu())
        if torch.cuda.is_available() and checkpoint_data['random_states'][3] is not None:
            torch.cuda.set_rng_state(checkpoint_data['random_states'][3].cpu())
        self.counters = checkpoint_data['counters']
        if 'losses' in checkpoint_data:  # Compatibility purposes until next release
            self.losses_epoch = checkpoint_data['losses']
        else:
            self.losses_epoch = checkpoint_data['losses_epoch']
            self.losses_it = checkpoint_data['losses_it']
        self.load_states_others(checkpoint_data) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def seed_all_rng(seed=None):
    """
    Set the random seed for the RNG in torch, numpy and python.

    Args:
        seed (int): if None, will use a strong random seed.
    """
    if seed is None:
        seed = (
            os.getpid()
            + int(datetime.now().strftime("%S%f"))
            + int.from_bytes(os.urandom(2), "big")
        )
        logger = logging.getLogger(__name__)
        logger.info("Using a generated random seed {}".format(seed))
    np.random.seed(seed)
    torch.set_rng_state(torch.manual_seed(seed).get_state())
    random.seed(seed)


# from https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def loadSeed(loadDir):
    """
    Loads the states and seed saved in a specified path
    
    Inputs:
        loadDir (path): where to look for thee seed to load; it is expected that
            the appropriate file within loadDir is named 'randomSeedUsed.pkl'
    
    Obs.: The file 'randomSeedUsed.pkl' should contain a list structured as
        follows. The length of this list is equal to the number of modules whose
        states were saved (torch, numpy, etc.). Each element in this list is a
        dictionary. The dictionary has three keys: 'module' with the name of 
        the module in string format ('numpy' or 'torch', for example), 'state' 
        with the saved generator state and, if corresponds, 'seed' with the 
        specific seed for the generator (note that torch has both state and 
        seed, but numpy only has state)
    """
    pathToSeed = os.path.join(loadDir, 'randomSeedUsed.pkl')
    with open(pathToSeed, 'rb') as seedFile:
        randomStates = pickle.load(seedFile)
        randomStates = randomStates['randomStates']
    for module in randomStates:
        thisModule = module['module']
        if thisModule == 'numpy':
            np.random.RandomState().set_state(module['state'])
        elif thisModule == 'torch':
            torch.set_rng_state(module['state'])
            torch.manual_seed(module['seed']) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def with_torch_seed(seed):
    assert isinstance(seed, int)
    rng_state = torch.get_rng_state()
    cuda_rng_state = torch.cuda.get_rng_state()
    set_torch_seed(seed)
    yield
    torch.set_rng_state(rng_state)
    torch.cuda.set_rng_state(cuda_rng_state) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def torch_seed(seed: Optional[int]):
    """Context manager which seeds the PyTorch PRNG with the specified seed and
    restores the state afterward. Setting seed to None is equivalent to running
    the code without the context manager."""
    if seed is None:
        yield
        return
    state = torch.get_rng_state()
    torch.manual_seed(seed)
    try:
        yield
    finally:
        torch.set_rng_state(state) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def write_dummy_file(filename, num_examples, maxlen):
    rng_state = torch.get_rng_state()
    torch.manual_seed(0)
    data = torch.rand(num_examples * maxlen)
    data = 97 + torch.floor(26 * data).int()
    with open(filename, "w") as h:
        offset = 0
        for _ in range(num_examples):
            ex_len = random.randint(1, maxlen)
            ex_str = " ".join(map(chr, data[offset : offset + ex_len]))
            print(ex_str, file=h)
            offset += ex_len
    torch.set_rng_state(rng_state) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def __exit__(self, type, value, traceback):
        torch.set_rng_state(self.old_rng_state)


# Conditional independence is recorded as a plate context at each site. 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def set_device_states(devices, states):
    for device, state in zip(devices, states):
        with torch.cuda.device(device):
            torch.cuda.set_rng_state(state) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def load_rng_state(file_name):
    rng_state = torch_load(file_name)
    torch.set_rng_state(rng_state) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_rng_state [as 别名]
def setstate(self, state: np.ndarray) -> None:
        return torch.set_rng_state(torch.from_numpy(state))