Stream

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

CUDA stream 的包裝函式。

CUDA stream 是一個屬於特定裝置的線性執行序列，獨立於其他 stream。請參閱 CUDA 語意 了解詳細資訊。

參數

• device ( torch.device 或 int, 可選) – 在其上分配 stream 的裝置。如果 device 為 None (預設) 或負整數，則將使用目前的裝置。

• priority ( int, 可選) – stream 的優先順序，應為 0 或負數，其中負數表示更高的優先順序。預設情況下，stream 的優先順序為 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 中目前排隊的核心：只有未來的操作會受到影響。