流

class torch.mtia.Stream(device, *, priority)

一個按序執行相關任務的非同步佇列，遵循先入先出（FIFO）原則。它可以控制或同步其他 Stream 的執行，或阻塞當前主機執行緒，以確保任務的正確排序。它支援使用 with 語句作為上下文管理器，確保 with 塊內的運算元在該流上執行。

有關適用於所有裝置的精確語義的詳細描述，請參閱對 CUDA 語義 行為的深入說明。

引數

• device (torch.device, 可選) – Stream 所需的裝置。如果未給出，將使用當前的 加速器 型別。

• priority (int, 可選) – 流的優先順序，應為 0 或負數，負數表示優先順序更高。預設情況下，流的優先順序為 0。

返回

一個 torch.Stream 物件。

返回型別

流

示例

>>> with torch.Stream(device='cuda') as s_cuda:
>>>     a = torch.randn(10, 5, device='cuda')
>>>     b = torch.randn(5, 10, device='cuda')
>>>     c = torch.mm(a, b)

query() → bool

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

返回

一個布林值，指示此流中的所有核心是否已完成。

返回型別

bool

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.query()
True

record_event(event) → Event

記錄一個事件。將其加入 Stream 佇列，以允許從 FIFO 佇列的當前點進行進一步同步。

引數

event (torch.Event, 可選) – 要記錄的事件。如果未給出，將分配一個新的事件。

返回

記錄的事件。

返回型別

Event

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> e_cuda = s_cuda.record_event()

synchronize() → None

等待此流中的所有核心完成。

示例

>>> s_cuda = torch.Stream(device='cuda')
>>> s_cuda.synchronize()

wait_event(event) → None

使提交到此流的所有後續工作等待某個事件。

引數

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

示例

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> e_cuda = s1_cuda.record_event()
>>> s2_cuda.wait_event(e_cuda)

wait_stream(stream) → None

與另一個流同步。所有提交到此流的後續工作將等待給定的流中所有已提交的核心完成後才開始執行。

引數

stream (torch.Stream) – 要同步的流。

示例

>>> s1_cuda = torch.Stream(device='cuda')
>>> s2_cuda = torch.Stream(device='cuda')
>>> s2_cuda.wait_stream(s1_cuda)