Stream

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

依照先進先出 (FIFO) 順序非同步執行各別任務的依序佇列。 它可以控制或同步其他 Stream 的執行，或封鎖目前的主機線程，以確保正確的任務排序。

請參閱 CUDA 語意 的深入說明，以了解適用於所有裝置的確切語意。

參數

• device (torch.device, optional) – Stream 的目標裝置。如果未指定，將使用目前的 加速器 類型。

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

傳回值

一個 torch.Stream 物件。

傳回類型

Stream

範例

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

query() → bool

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

傳回值

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

傳回類型

bool

範例

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

record_event(event) → Event

記錄一個事件。 將其排隊到 Stream 中，以便從 FIFO 隊列中的當前點進行進一步的同步。

參數

event (torch.Event, optional) – 要記錄的事件。 如果未指定，將分配一個新的事件。

傳回值

已記錄的事件。

傳回類型

Event

範例

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

synchronize() → None

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

範例

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

wait_event(event) → None

使所有未來提交到 Stream 的工作等待一個事件。

參數

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 同步。所有未來提交到此 Stream 的工作將等待直到已提交到給定 Stream 的所有核心完成。

參數

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

範例

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