流

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 中當前入隊的核心的情況下返回：只有未來的操作會受到影響。