torch.cuda.make_graphed_callables

torch.cuda.make_graphed_callables(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)

接受可呼叫物件 (函數或 nn.Module) 並傳回繪圖版本。

每個繪圖可呼叫物件的 forward pass 都會將其來源可呼叫物件的 forward CUDA 工作，以 CUDA 圖形的形式在單一 autograd 節點內執行。

繪圖可呼叫物件的 forward pass 也會將一個 backward 節點附加到 autograd 圖形。在 backward 期間，此節點會將可呼叫物件的 backward 工作以 CUDA 圖形的形式執行。

因此，每個繪圖可呼叫物件都應該可以即插即用地取代其來源可呼叫物件，並用於啟用 autograd 的訓練迴圈中。

有關詳細的使用方式和約束，請參閱 Partial-network capture。

如果您傳遞多個可呼叫物件的元組，它們的捕捉將使用相同的記憶體池。有關何時適合使用這種方式，請參閱 Graph memory management。

參數

• callables ( torch.nn.Module 或 Python 函數，或 tuple of these) – 要繪圖的可呼叫物件或多個可呼叫物件。有關何時適合傳遞可呼叫物件的元組，請參閱 Graph memory management。如果您傳遞可呼叫物件的元組，則它們在元組中的順序必須與它們在即時工作負載中執行的順序相同。

• sample_args ( tuple of Tensors, or tuple of tuples of Tensors) – 每個可呼叫物件的範例引數。如果傳遞單個可呼叫物件，則 sample_args 必須是引數 Tensor 的單個元組。如果傳遞可呼叫物件的元組，則 sample_args 必須是引數 Tensor 的元組的元組。

• num_warmup_iters ( int) – 熱身迭代的次數。目前，DataDistributedParallel 需要 11 次迭代才能熱身。預設值：3。

• allow_unused_input ( bool) – 如果為 False，則指定在計算輸出時未使用的輸入 (因此它們的 grad 始終為零) 將會出錯。預設為 False。

• pool (optional) – Token (由 graph_pool_handle() 或 other_Graph_instance.pool() 傳回) 提示此圖形可能會與指定的池共用記憶體。請參閱 Graph memory management。

注意

sample_args 中每個 Tensor 的 requires_grad 狀態必須與訓練迴圈中對應實際輸入所期望的狀態相符。

警告

此 API 處於 Beta 階段，未來版本可能會變更。

警告

每個可呼叫物件的 sample_args 必須僅包含 Tensor。不允許其他類型。

警告

傳回的可呼叫物件不支援高階微分 (例如，double backward)。

警告

在傳遞給 make_graphed_callables() 的任何 Module 中，只有參數可以是可訓練的。緩衝區必須具有 requires_grad=False。

警告

在您通過 make_graphed_callables() 傳遞 torch.nn.Module 之後，您不得新增或移除該 Module 的任何參數或緩衝區。

警告

傳遞給 make_graphed_callables() 的 torch.nn.Module 在傳遞時不得在其上註冊模組 hook。但是，允許在通過 make_graphed_callables() 之後 在模組上註冊 hook。

警告

執行繪圖可呼叫物件時，您必須以其 sample_args 中出現的相同順序和格式傳遞其引數。

警告

只有在停用快取的情況下，make_graphed_callables() 才支援自動混合精度。context manager torch.cuda.amp.autocast() 必須具有 cache_enabled=False。