Функция cuda.set_rng_state_all

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

Синтаксис

torch.cuda.set_rng_state_all(state)

Пример

Давайте сохраним состояние генераторов случайных чисел на всех GPU, сгенерируем тензоры на каждом устройстве, а затем восстановим состояние и убедимся, что последовательности случайных чисел повторяются:

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	device_count = torch.cuda.device_count()
	print(f"Доступно устройств: {device_count}")
	
	# Сохраняем состояние на всех устройствах
	state_all = torch.cuda.get_rng_state_all()
	
	# Генерируем тензоры на каждом устройстве
	tensors1 = []
	for i in range(device_count):
		t = torch.rand(3, device=f'cuda:{i}')
		tensors1.append(t)
	
	# Восстанавливаем состояние на всех устройствах
	torch.cuda.set_rng_state_all(state_all)
	
	# Генерируем тензоры снова
	tensors2 = []
	for i in range(device_count):
		t = torch.rand(3, device=f'cuda:{i}')
		tensors2.append(t)
	
	# Сравниваем результаты
	for i in range(device_count):
		print(f"Устройство {i}: {torch.all(tensors1[i] == tensors2[i])}")

Результат выполнения кода (пример для двух устройств):

Доступно устройств: 2
Устройство 0: True
Устройство 1: True

Пример

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

import torch

torch.manual_seed(42)
if torch.cuda.is_available():
	device_count = torch.cuda.device_count()
	
	# Сохраняем текущее состояние
	state_all = torch.cuda.get_rng_state_all()
	
	# Генерируем тензор на первом устройстве
	t1 = torch.rand(5, device='cuda:0')
	
	# Восстанавливаем состояние
	torch.cuda.set_rng_state_all(state_all)
	
	# Генерируем тензор снова
	t2 = torch.rand(5, device='cuda:0')
	
	# Проверяем, что тензоры одинаковы
	print(f"Тензоры равны: {torch.all(t1 == t2)}")
	print(f"Первый тензор: {t1}")
	print(f"Второй тензор: {t2}")

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

Тензоры равны: True
Первый тензор: tensor([0.8823, 0.9150, 0.3829, 0.9593, 0.3904], device='cuda:0')
Второй тензор: tensor([0.8823, 0.9150, 0.3829, 0.9593, 0.3904], device='cuda:0')

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

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