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）並返回圖化版本。

每個圖化可呼叫物件的前向傳播將其源可呼叫物件的前向 CUDA 工作作為 CUDA 圖在一個 autograd 節點內執行。

圖化可呼叫物件的前向傳播還會向 autograd 圖新增一個反向傳播節點。在反向傳播期間，該節點將其可呼叫物件的反向傳播工作作為 CUDA 圖執行。

因此，在啟用 autograd 的訓練迴圈中，每個圖化可呼叫物件都應是其源可呼叫物件的直接替代品。

有關詳細用法和限制，請參閱 部分網路捕獲。

如果您傳遞多個可呼叫物件的元組，它們的捕獲將使用相同的記憶體池。有關何時適用此情況，請參閱 圖記憶體管理。

引數

• callables (torch.nn.Module 或 Python 函式，或 元組 of these) – 要圖化（graph）的可呼叫物件或可呼叫物件列表。有關何時適合傳遞可呼叫物件的元組，請參閱 圖記憶體管理。如果您傳遞可呼叫物件的元組，它們在元組中的順序必須與它們在實際工作負載中執行的順序相同。

• sample_args (Tensor 的 元組，或 Tensor 的 元組 的 元組) – 每個可呼叫物件的示例引數。如果傳遞了單個可呼叫物件，則 sample_args 必須是引數 Tensor 的單個元組。如果傳遞了可呼叫物件的元組，則 sample_args 必須是引數 Tensor 的元組的元組。

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

• allow_unused_input (bool) – 如果為 False，則指定在計算輸出時未使用（因此其梯度始終為零）的輸入將導致錯誤。預設值為 False。

• pool (可選) – 令牌（由 graph_pool_handle() 或 other_Graph_instance.pool() 返回），提示此圖可能與指定的記憶體池共享記憶體。請參閱 圖記憶體管理。

注意

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

警告

此 API 處於 beta 階段，未來版本中可能會發生變化。

警告

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

警告

返回的可呼叫物件不支援高階微分（例如，二次反向傳播）。

警告

傳遞給 make_graphed_callables() 的任何 Module 中，只有引數可以是可訓練的。Buffer 必須設定 requires_grad=False。

警告

將 torch.nn.Module 傳遞給 make_graphed_callables() 後，您不能新增或移除該 Module 的任何引數或 buffer。

警告

傳遞給 make_graphed_callables() 的 torch.nn.Module 在傳遞時不得註冊任何模組鉤子。但是，在將模組透過 make_graphed_callables() *之後*註冊鉤子是允許的。

警告

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

警告

自動混合精度在 make_graphed_callables() 中僅在停用快取時受支援。上下文管理器 torch.cuda.amp.autocast() 必須設定 cache_enabled=False。