流

classtorch.Stream(device, *, priority)

一个按先入先出（FIFO）顺序异步执行任务的队列。它能控制和同步其他Stream中的任务执行，并通过阻塞当前线程来保证任务的正确顺序。

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

参数

• device (torch.device, 可选) – 流的期望设备。如果没有指定，则使用当前的 加速器 类型。

• priority (int, 可选) – 流的优先级，默认值为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 参数并返回 None。

与此流同步。提交到该流的所有未来任务将等待之前提交到该流的所有内核完成。

参数

stream (torch.Stream) - 需要同步的流。

示例：

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