流

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

CUDA 流的封装。

CUDA 流是属于特定设备的一系列线性执行操作，与其他流无关。它支持 `with` 语句作为上下文管理器，以确保 `with` 块内的操作在该流上运行。有关详细信息，请参阅 CUDA 语义。

参数

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

• priority (int, 可选) – 流的优先级，可以是正数、0 或负数。数字越小表示优先级越高。默认优先级为 0。如果值超出允许的优先级范围，它将自动映射到最近的有效优先级（正数过大则映射到最低优先级，负数过大则映射到最高优先级）。

query()

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

返回

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

返回类型

布尔值

record_event(event=None)

Record an event.

参数

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

返回

记录的事件。

synchronize()

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

注意

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

wait_event(event)

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

参数

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

注意

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

此函数在不等待 event 的情况下返回：只有未来的操作会受到影响。

wait_stream(stream)

与另一个流同步。

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

参数

stream (Stream) – 要同步的流。

注意

此函数在不等待 stream 中当前入队的内核的情况下返回：只有未来的操作会受到影响。