[docs]defget_rng_state(device:Union[int,str,torch.device]="cuda")->Tensor:r"""Return the random number generator state of the specified GPU as a ByteTensor. Args: device (torch.device or int, optional): The device to return the RNG state of. Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device). .. warning:: This function eagerly initializes CUDA. """_lazy_init()ifisinstance(device,str):device=torch.device(device)elifisinstance(device,int):device=torch.device("cuda",device)idx=device.indexifidxisNone:idx=current_device()default_generator=torch.cuda.default_generators[idx]returndefault_generator.get_state()
[docs]defget_rng_state_all()->List[Tensor]:r"""Return a list of ByteTensor representing the random number states of all devices."""results=[get_rng_state(i)foriinrange(device_count())]returnresults
[docs]defset_rng_state(new_state:Tensor,device:Union[int,str,torch.device]="cuda")->None:r"""Set the random number generator state of the specified GPU. Args: new_state (torch.ByteTensor): The desired state device (torch.device or int, optional): The device to set the RNG state. Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device). """withtorch._C._DisableFuncTorch():new_state_copy=new_state.clone(memory_format=torch.contiguous_format)ifisinstance(device,str):device=torch.device(device)elifisinstance(device,int):device=torch.device("cuda",device)defcb():idx=device.indexifidxisNone:idx=current_device()default_generator=torch.cuda.default_generators[idx]default_generator.set_state(new_state_copy)_lazy_call(cb)
[docs]defset_rng_state_all(new_states:Iterable[Tensor])->None:r"""Set the random number generator state of all devices. Args: new_states (Iterable of torch.ByteTensor): The desired state for each device. """fori,stateinenumerate(new_states):set_rng_state(state,i)
[docs]defmanual_seed(seed:int)->None:r"""Set the seed for generating random numbers for the current GPU. It's safe to call this function if CUDA is not available; in that case, it is silently ignored. Args: seed (int): The desired seed. .. warning:: If you are working with a multi-GPU model, this function is insufficient to get determinism. To seed all GPUs, use :func:`manual_seed_all`. """seed=int(seed)defcb():idx=current_device()default_generator=torch.cuda.default_generators[idx]default_generator.manual_seed(seed)_lazy_call(cb,seed=True)
[docs]defmanual_seed_all(seed:int)->None:r"""Set the seed for generating random numbers on all GPUs. It's safe to call this function if CUDA is not available; in that case, it is silently ignored. Args: seed (int): The desired seed. """seed=int(seed)defcb():foriinrange(device_count()):default_generator=torch.cuda.default_generators[i]default_generator.manual_seed(seed)_lazy_call(cb,seed_all=True)
[docs]defseed()->None:r"""Set the seed for generating random numbers to a random number for the current GPU. It's safe to call this function if CUDA is not available; in that case, it is silently ignored. .. warning:: If you are working with a multi-GPU model, this function will only initialize the seed on one GPU. To initialize all GPUs, use :func:`seed_all`. """defcb():idx=current_device()default_generator=torch.cuda.default_generators[idx]default_generator.seed()_lazy_call(cb)
[docs]defseed_all()->None:r"""Set the seed for generating random numbers to a random number on all GPUs. It's safe to call this function if CUDA is not available; in that case, it is silently ignored. """defcb():random_seed=0seeded=Falseforiinrange(device_count()):default_generator=torch.cuda.default_generators[i]ifnotseeded:default_generator.seed()random_seed=default_generator.initial_seed()seeded=Trueelse:default_generator.manual_seed(random_seed)_lazy_call(cb)
[docs]definitial_seed()->int:r"""Return the current random seed of the current GPU. .. warning:: This function eagerly initializes CUDA. """_lazy_init()idx=current_device()default_generator=torch.cuda.default_generators[idx]returndefault_generator.initial_seed()
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.