流 ¶

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

CUDA 流包装器。

CUDA 流是特定设备上的线性执行序列，与其他流独立。它支持 with 语句作为上下文管理器，以确保 with 块内的操作在相应的流上运行。有关 CUDA 语义的详细信息，请参阅 CUDA 文档。

参数:

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

• priority (int，可选) – 流的优先级，可以是正数、0 或负数。数字越小表示优先级越高。默认情况下，优先级设置为 0。如果值超出允许的优先级范围，它将自动映射到最近的合法优先级（对于大正数是最低优先级，对于大负数是最高优先级）。

query()

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

返回值:

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

返回类型:

布尔型

record_event(event=None)[source][source]

记录事件。

参数:

event (torch.cuda.Event, 可选) – 要记录的事件。如果未提供，将分配一个新的。

返回值:

已记录的事件。

synchronize()

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

注意

这是一个围绕 cudaStreamSynchronize() 的包装器：请参阅 CUDA Stream 文档以获取更多信息。

wait_event(事件)[源代码][源代码] ¶

使所有提交给流的后续工作等待事件。

参数:

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

注意

这是一个围绕 cudaStreamWaitEvent() 的包装器：请参阅 CUDA Stream 文档以获取更多信息。

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

wait_stream(流)[源代码][源代码] ¶

与另一个流同步。

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

参数:

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

注意

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