流

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

一个按顺序执行相应任务的异步队列，遵循先进先出 (FIFO) 顺序。它可以控制或同步其他流的执行，或阻塞当前主机线程以确保正确的任务排序。

有关适用于所有设备的确切语义的详细信息，请参阅 CUDA 语义 中对 CUDA 行为的深入描述。

参数

• device (torch.device，可选) – 流所需的设备。如果未给出，则将使用当前 加速器 类型。

• priority (int, 可选) – 流的优先级，应为 0 或负数，其中负数表示更高的优先级。默认情况下，流的优先级为 0。

返回值

一个 torch.Stream 对象。

返回类型

流

示例

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

query() → bool

检查是否已完成提交的所有工作。

返回值

一个布尔值，指示此流中的所有内核是否已完成。

返回类型

bool

示例

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

record_event(event) → Event

记录事件。将其排队到流中，以便从 FIFO 队列中的当前位置进行进一步同步。

参数

event (torch.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)