torch.compiler.allow_in_graph

torch.compiler.allow_in_graph(fn)

告知編譯器前端 (Dynamo) 跳過函式的符號化內省，而是在遇到時直接將其寫入圖中。

如果您正在使用 torch.compile() (搭配 backend=”inductor” (預設值))，或 torch.export.export()，並且嘗試在所有追蹤過程中將一個 Python 函式視為黑箱，請不要使用此 API。請改為建立一個自訂運算子（請參閱 PyTorch 自訂運算子登陸頁面）。

警告

如果您是典型的 torch.compile 使用者 (例如，您正在將 torch.compile 應用於模型以使其運行更快)，您可能不想使用此函式。allow_in_graph() 是一個潛在的錯誤來源，因為它會跳過負責執行安全檢查 (圖形中斷、處理閉包等) 的編譯器前端 (Dynamo)。不正確的使用將導致難以除錯的靜默不正確性問題。

給定一個沒有 allow_in_graph 裝飾器的 Python 函式，torch.compile 的常規執行會追蹤該函式。allow_in_graph() 改變了它，使前端不會在函式內部追蹤，但編譯器後端仍然會追蹤它。將其與自訂運算子進行比較，自訂運算子將函式視為整個 torch.compile 堆疊中的黑箱。下表比較了這些機制。

機制	前端 (Dynamo)	後端 (AOTAutograd+Inductor)
無裝飾器	追蹤內部	追蹤內部
允許在圖中 (allow_in_graph)	不透明的可呼叫物件	追蹤內部
自訂運算子	不透明的可呼叫物件	不透明的可呼叫物件

allow_in_graph() 的一個常見用例是作為編譯器前端的逃生艙：如果您知道該函式適用於編譯堆疊的下游組件 (AOTAutograd 和 Inductor)，但 Dynamo 存在一個錯誤，阻止它正確地以符號方式內省該函式 (或者如果您的程式碼是用 C/C++ 編寫的，因此無法使用 Dynamo 進行內省)，那麼可以使用 allow_in_graph() 裝飾該函式來繞過 Dynamo。

我們要求 fn 遵守以下限制。未能遵守會導致未定義的行為

• fn 的輸入必須是 FX 圖中可代理的類型。有效類型包括：Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?] Tuple[Tensor?, …]/Tuple[int?, …]/Tuple[float?, …]/torch.dtype/torch.device

• fn 的輸出必須是 FX 圖中可代理的類型 (請參閱前一個項目)

• fn 內部使用的所有 Tensor 必須直接作為輸入傳遞到 fn (而不是作為捕獲的變數)。

參數

fn – 代表要包含在圖形中的函式的可呼叫物件。如果 fn 是一個可呼叫物件的列表或元組，它會遞迴地將 allow_in_graph() 應用於每個函式，並傳回包含修改後函式的新列表或元組。

範例

torch.compiler.allow_in_graph(my_custom_function)

@torch.compile(...)
def fn(x):
    x = torch.add(x, 1)
    x = my_custom_function(x)
    x = torch.add(x, 1)
    return x

fn(...)

將捕獲包含 my_custom_function() 的單個圖。