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 中當前已入隊的核函式：只有未來的操作會受到影響。