Функция cuda.get_rng_state_all

Функция cuda.get_rng_state_all возвращает кортеж, содержащий состояния генераторов случайных чисел для всех доступных устройств CUDA. Это полезно, когда необходимо сохранить полное состояние случайности на всех GPU для последующего восстановления с помощью cuda.set_rng_state_all. Функция не принимает параметров.

Синтаксис

torch.cuda.get_rng_state_all()

Пример

Давайте получим состояния генераторов для всех доступных GPU:

import torch

if torch.cuda.is_available():
	states = torch.cuda.get_rng_state_all()
	print(len(states))
else:
	print("CUDA не доступна")

Результат выполнения кода:

"2"

В результате мы получили кортеж из двух состояний для двух устройств CUDA.

Пример

Давайте сохраним состояния генераторов, создадим тензоры на всех устройствах, а затем восстановим состояния:

import torch

torch.manual_seed(0)

if torch.cuda.is_available():
	# Сохраняем состояния
	saved_states = torch.cuda.get_rng_state_all()
	
	# Генерируем тензоры на каждом устройстве
	for i in range(torch.cuda.device_count()):
		with torch.cuda.device(i):
			t = torch.rand(3, device='cuda')
			print(f'Device {i}: {t}')
	
	# Восстанавливаем состояния
	torch.cuda.set_rng_state_all(saved_states)
	
	# Генерируем снова - должны получить те же значения
	for i in range(torch.cuda.device_count()):
		with torch.cuda.device(i):
			t = torch.rand(3, device='cuda')
			print(f'Restored device {i}: {t}')
else:
	print("CUDA не доступна")

Результат выполнения кода:

"Device 0: tensor([0.4963, 0.7682, 0.0885], device='cuda:0')"
"Device 1: tensor([0.1320, 0.3074, 0.6341], device='cuda:1')"
"Restored device 0: tensor([0.4963, 0.7682, 0.0885], device='cuda:0')"
"Restored device 1: tensor([0.1320, 0.3074, 0.6341], device='cuda:1')"

Как видно из примера, после восстановления состояний сгенерированные тензоры полностью совпадают с первоначальными.

Пример

Давайте сравним состояния отдельных устройств и всех устройств сразу:

import torch

torch.manual_seed(42)

if torch.cuda.is_available():
	# Получаем состояния всех устройств
	all_states = torch.cuda.get_rng_state_all()
	
	# Получаем состояние конкретного устройства
	single_state = torch.cuda.get_rng_state(0)
	
	# Сравниваем первое состояние кортежа с состоянием устройства 0
	is_equal = torch.equal(all_states[0], single_state)
	print(is_equal)
else:
	print("CUDA не доступна")

Результат выполнения кода:

"True"

Первое состояние кортежа соответствует состоянию генератора на устройстве 0.

Смотрите также

функцию cuda.set_rng_state_all,
которая восстанавливает состояния генераторов на всех устройствах CUDA
функцию cuda.get_rng_state,
которая возвращает состояние генератора для конкретного устройства
функцию get_rng_state,
которая возвращает состояние генератора для CPU
функцию cuda.manual_seed_all,
которая устанавливает начальное зерно для всех устройств CUDA