Stream

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

CUDA stream 的包装器。

CUDA stream 是属于特定设备的一条线性执行序列，独立于其他 stream。它支持使用 with 语句作为上下文管理器，以确保 with 块内的算子在该 stream 上运行。详见 CUDA semantics。

参数

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

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

query()

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

返回

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

返回类型

bool

record_event(event=None)

记录一个事件。

参数

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

返回

记录的事件。

synchronize()

等待此 stream 中的所有核函数完成。

注意

这是 cudaStreamSynchronize() 的包装器：详见 CUDA Stream 文档 获取更多信息。

wait_event(event)

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

参数

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

注意

这是 cudaStreamWaitEvent() 的包装器：详见 CUDA Stream 文档 获取更多信息。

此函数会立即返回，不会等待 event 完成：只有未来的操作会受到影响。

wait_stream(stream)

与另一个 stream 同步。

所有提交到此 stream 的未来工作将等到调用时已提交到给定 stream 的所有核函数完成后再执行。

参数

stream (Stream) – 要同步的 stream。

注意

此函数会立即返回，不会等待 stream 中当前已入队的核函数：只有未来的操作会受到影响。