functorch.compile.aot_module

functorch.compile.aot_module(mod, *args, **kwargs)

使用 torch 調度追蹤機制追蹤 mod 的正向和反向圖。它是一個包裝函式，底層使用 aot_function() 執行追蹤和編譯。

aot_module() 會將 nn.Module 的參數和緩衝區提升為一個新的可調呼叫物件的輸入，然後透過 aot_function() 編譯。

警告

此 API 仍處於實驗階段，可能會有所變更。

參數

• mod (Callable) – 一個 nn.Module 模組。

• args – 要傳遞給 aot_function() 的參數

• kwargs – 要傳遞給 aot_function() 的關鍵字參數

回傳值

回傳一個 nn.Module，它保留了原始 mod 的 Eager 模式行為，但其正向和反向圖已被編譯。