流 ¶

class torch.Stream(设备, *, 优先级)

按照先进先出（FIFO）顺序异步执行各自任务的有序队列。它可以控制或同步其他 Stream 的执行，或阻塞当前主机线程以确保正确的任务顺序。它支持 with 语句作为上下文管理器，以确保 with 块内的操作在相应的 Stream 上运行。

详细了解 CUDA 行为，请参阅 CUDA 语义中的 CUDA 行为描述，以了解适用于所有设备的确切语义。

参数:

• device（ torch.device ，可选）- Stream 期望的设备。如果没有提供，将使用当前加速器类型。

• priority（int，可选）- Stream 的优先级，应为 0 或负数，其中负数表示更高的优先级。默认情况下，Stream 的优先级为 0。

返回值:

一个 torch.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)

query() → bool

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

返回值:

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

返回类型:

布尔型

示例：

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

record_event(event) → 事件 ¶

记录一个事件。将其入队到 Stream 中，以便从当前 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

与另一个流同步。所有提交给此流的后续工作将等待所有已提交给给定流的内核完成。

参数:

流（ torch.Stream ）- 用于同步的流。

示例：

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