Python torch.set_rng_state() Examples

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 

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) 

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) 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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 

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) 

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) 

def setUp(self) -> None:
        torch.set_rng_state(torch.manual_seed(42).get_state()) 

def _load_metadata(self, checkpoint):
        np.random.set_state(checkpoint[NP_RANDOM_STATE])
        random.setstate(checkpoint[PYTHON_RANDOM_STATE])
        torch.set_rng_state(checkpoint[PYTORCH_RANDOM_STATE]) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def _load_metadata(self, checkpoint):
        np.random.set_state(checkpoint[NP_RANDOM_STATE])
        random.setstate(checkpoint[PYTHON_RANDOM_STATE])
        torch.set_rng_state(checkpoint[PYTORCH_RANDOM_STATE]) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def setUp(self) -> None:
        torch.set_rng_state(torch.manual_seed(42).get_state()) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state) 

def tearDown(self):
        if hasattr(self, "rng_state"):
            torch.set_rng_state(self.rng_state)