Поток

class torch.mtia.Stream(device, *, priority)

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

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

Параметры:

• device (torch.device, необязательно) – требуемое устройство для потока. Если не задано, будет использоваться текущий тип ускорителя.
• 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

Записывает событие. Добавляет его в очередь потока, чтобы обеспечить дальнейшую синхронизацию с текущей точки очереди 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)