ExternalStream

class torch.cuda.ExternalStream(stream_ptr, device=None, **kwargs)

包裝一個外部分配的 CUDA 流。

此類用於包裝在其他庫中分配的流，以方便資料交換和多庫互動。

注意

此類不管理流的生命週期，使用者在使用此類時有責任確保引用的流保持活動狀態。

引數

• stream_ptr (int) – 外部分配的 cudaStream_t 值的整數表示。

• device (torch.device 或 int, 可選) – 流最初分配的裝置。如果裝置指定不正確，則使用此流的後續啟動可能會失敗。

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 中當前排隊的核心的情況下返回：僅影響未來的操作。