Stream¶

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

一個按照先進先出 (FIFO) 順序非同步執行相應任務的有序佇列。它可以控制或同步其他 Stream 的執行，或者阻塞當前主機執行緒以確保任務的正確順序。它支援 with 語句作為上下文管理器，以確保 with 程式碼塊中的運算子在該 Stream 上執行。

有關適用於所有裝置的精確語義的詳細資訊，請參閱 CUDA 語義的深入描述。

引數

device (torch.device, 可選) – Stream 所需的裝置。如果未提供，將使用當前的加速器型別。
priority (int, 可選) – Stream 的優先順序，應為 0 或負數，其中負數表示更高的優先順序。預設情況下，Stream 的優先順序為 0。

返回

一個 torch.Stream 物件。

返回型別

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¶

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

返回: 一個布林值，指示此 Stream 中的所有核心是否已完成。
返回型別: 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¶

等待此 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)

Stream¶

文件

教程

資源