Stream

class torch.Stream(device, *, priority)

Упорядоченная очередь, в которой соответствующие задачи выполняются асинхронно в порядке «первым поступил — первым выполнен» (FIFO). Она позволяет управлять выполнением других потоков Stream или синхронизировать его, а также блокировать текущий поток хоста, чтобы обеспечить правильную последовательность выполнения задач. Stream поддерживает оператор with в качестве менеджера контекста, обеспечивая выполнение операторов внутри блока with в соответствующем потоке.

Подробное описание поведения CUDA см. в разделе Семантика CUDA, где описана точная семантика, применяемая ко всем устройствам.

Параметры:

• device (torch.device, необязательно) – требуемое устройство для Stream. Если не указано, будет использоваться текущий тип ускорителя.
• priority (int, необязательно) – приоритет потока; должен быть равен 0 или быть отрицательным. Отрицательные значения обозначают более высокий приоритет. По умолчанию приоритет потоков равен 0.

Возвращает:

Объект torch.Stream.

Тип возвращаемого значения:

Stream

Пример:

>>> with torch.Stream(device='cuda') as s_cuda:
>>>     a = torch.randn(10, 5, device='cuda')
>>>     b = torch.randn(5, 10, device='cuda')
>>>     c = torch.mm(a, b)

is_capturing() → bool

Возвращает true, если в этом потоке в данный момент записываются операции для захвата графа.

Возвращает:

Логическое значение, указывающее, выполняется ли захват в потоке.

Тип возвращаемого значения:

bool

Пример:

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.is_capturing()

query() → bool

Проверяет, завершена ли вся отправленная работа.

Возвращает:

Логическое значение, указывающее, завершены ли все ядра в этом потоке.

Тип возвращаемого значения:

bool

Пример:

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.query()
True

record_event(event) → Event

Записывает событие. Добавляет его в очередь Stream, чтобы обеспечить последующую синхронизацию с текущей позиции в очереди FIFO.

Параметры:

event (torch.Event, необязательно) – событие для записи. Если не указано, будет создано новое.

Возвращает:

Записанное событие.

Тип возвращаемого значения:

Event

Пример:

>>> s_cuda = torch.Stream(device='cuda')
>>> e_cuda = s_cuda.record_event()

synchronize() → None

Ожидает завершения всех ядер в этом потоке.

Пример:

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.synchronize()

wait_event(event) → None

Заставляет всю последующую работу, отправленную в поток, ожидать наступления события.

Параметры:

event (torch.Event) – событие, наступления которого нужно дождаться.

Пример:

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> e_cuda = s1_cuda.record_event()
>>> s2_cuda.wait_event(e_cuda)

wait_stream(stream) → None

Синхронизируется с другим потоком. Вся последующая работа, отправленная в этот поток, будет ожидать завершения всех ядер, ранее отправленных в указанный поток.

Параметры:

stream (torch.Stream) – поток для синхронизации.

Пример:

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> s2_cuda.wait_stream(s1_cuda)