Класс Stream

Класс Stream в PyTorch предоставляет интерфейс для работы с потоками CUDA. Потоки позволяют выполнять операции на GPU асинхронно и параллельно, что значительно повышает производительность при работе с большими тензорами. Основное назначение класса - управление очередями выполнения операций на устройстве CUDA.

Синтаксис

torch.cuda.Stream(device=-1, priority=0)

Параметры конструктора:

device (int, опционально) - идентификатор устройства CUDA. По умолчанию используется текущее устройство.
priority (int, опционально) - приоритет потока. Может принимать значения от -1 (высокий) до 0 (низкий), поддерживается не всеми устройствами.

Пример

Создадим поток и выполним в нём операции с тензорами:

import torch

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

# Создаём тензоры на GPU
t1 = torch.tensor([1, 2, 3, 4, 5], device='cuda')
t2 = torch.tensor([5, 4, 3, 2, 1], device='cuda')

# Выполняем операцию в потоке
with torch.cuda.stream(stream):
	res = t1 + t2
	stream.synchronize()

print(res)

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

tensor([6, 6, 6, 6, 6], device='cuda:0')

Пример

Покажем неблокирующее выполнение операций в разных потоках:

import torch
import time

torch.manual_seed(0)
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()

t1 = torch.randn(1000, 1000, device='cuda')
t2 = torch.randn(1000, 1000, device='cuda')

# Запускаем операции в разных потоках
with torch.cuda.stream(stream1):
	res1 = t1 * 2
	stream1.synchronize()

with torch.cuda.stream(stream2):
	res2 = t2 * 3
	stream2.synchronize()

print(res1.shape, res2.shape)

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

torch.Size([1000, 1000]) torch.Size([1000, 1000])

Пример

Использование потока в качестве контекстного менеджера:

import torch

stream = torch.cuda.Stream()

# Создаём тензоры
t1 = torch.tensor([1, 2, 3, 4, 5], device='cuda')
t2 = torch.tensor([5, 4, 3, 2, 1], device='cuda')

# Входим в контекст потока
with torch.cuda.stream(stream):
	res = t1 * t2

# Ожидаем завершения всех операций в потоке
stream.synchronize()
print(res)

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

tensor([5, 8, 9, 8, 5], device='cuda:0')

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

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