Метод synchronize

Метод synchronize класса Stream вызывает блокировку текущего потока выполнения до тех пор, пока все операции, запущенные в указанном потоке, не будут полностью завершены. Это гарантирует, что после вызова данного метода все вычисления в целевом потоке завершены, и их результаты можно безопасно использовать в основном потоке или других потоках. Метод не принимает никаких параметров и не возвращает значений.

Синтаксис

stream.synchronize()

Пример

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

import torch

# Создаем поток CUDA
stream = torch.cuda.Stream()

# Выполняем операции в потоке
with torch.cuda.stream(stream):
	t = torch.tensor([1, 2, 3, 4, 5], device='cuda')
	res = t * 2

# Синхронизируем поток
stream.synchronize()

# Теперь можно безопасно использовать результат
print(res.cpu())

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

tensor([ 2, 4, 6, 8, 10])

Пример

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

import torch

torch.manual_seed(0)

# Создаем два потока
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

# Операции в первом потоке
with torch.cuda.stream(s1):
	t1 = torch.randn(1000, device='cuda')
	t1 = t1 * 2

# Синхронизируем первый поток
s1.synchronize()

# Операции во втором потоке используют результат первого
with torch.cuda.stream(s2):
	t2 = t1 + 1
	t2 = t2.sum()

# Синхронизируем второй поток
s2.synchronize()

print(t2.cpu())

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

tensor(972.2628)

Пример

Важно помнить, что метод synchronize блокирует текущий поток до полного завершения всех операций в целевом потоке. Это может быть полезно при отладке или для гарантии корректности данных:

import torch

# Создаем поток для длительных вычислений
stream = torch.cuda.Stream()

with torch.cuda.stream(stream):
	t = torch.zeros(1000000, device='cuda')
	for i in range(100):
		t += i

# Проверяем статус выполнения
print(f"До синхронизации: {stream.query()}")

# Блокируем поток до завершения
stream.synchronize()

print(f"После синхронизации: {stream.query()}")
print(f"Результат: {t.sum().cpu()}")

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

До синхронизации: False
После синхронизации: True
Результат: 49500000.0

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

класс Stream,
который представляет поток CUDA для параллельного выполнения операций
метод wait_stream,
который синхронизирует текущий поток с другим потоком
метод record_event,
который записывает событие в поток
метод query,
который проверяет, завершились ли все операции в потоке