Stream

class torch.cuda.Stream(device=None, priority=0, **kwargs)

CUDA 流的包装器。

CUDA 流是属于特定设备的线性执行序列，独立于其他流。有关详细信息，请参阅 CUDA 语义。

参数

• device (torch.device or int, optional) – 在其上分配流的设备。如果 device 为 None（默认）或负整数，则将使用当前设备。

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

query()

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

返回值

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

返回类型

bool

record_event(event=None)

记录一个事件。

参数

event (torch.cuda.Event, optional) – 要记录的事件。如果未给出，将分配一个新的事件。

返回值

已记录的事件。

synchronize()

等待此流中的所有内核完成。

注意

这是 cudaStreamSynchronize() 的包装器：有关更多信息，请参阅 CUDA Stream 文档。

wait_event(event)

使提交到流的所有未来工作等待一个事件。

参数

event (torch.cuda.Event) – 要等待的事件。

注意

这是 cudaStreamWaitEvent() 的包装器：有关更多信息，请参阅 CUDA Stream 文档。

此函数返回而不等待事件：仅影响未来的操作。

wait_stream(stream)

与另一个流同步。

提交到此流的所有未来工作将等待，直到在调用时提交到给定流的所有内核完成。

参数

stream (Stream) – 要同步的流。

注意

此函数返回而不等待 stream 中当前排队的内核：仅影响未来的操作。