評價此頁
★ ★ ★ ★ ★

torch.compile#

torch.compile(model: Callable[[_InputT], _RetT], *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool, Callable]]] = None, disable: bool = False) → Callable[[_InputT], _RetT][原始碼]#
torch.compile(model: None = None, *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool, Callable]]] = None, disable: bool = False) → Callable[[Callable[[_InputT], _RetT]], Callable[[_InputT], _RetT]]

使用 TorchDynamo 和指定的後端最佳化給定的模型/函式。如果您正在編譯一個 torch.nn.Module，您還可以使用 torch.nn.Module.compile() 原地編譯模組，而無需更改其結構。

具體來說，對於在編譯區域內執行的每個幀，我們將嘗試編譯它，並將編譯後的結果快取在程式碼物件中以供將來使用。單個幀可能會被編譯多次，如果先前的編譯結果不適用於後續呼叫（這稱為“guard failure”），您可以使用 TORCH_LOGS=guards 來除錯這些情況。最多可以將多個編譯結果與一個幀關聯起來，直到 torch._dynamo.config.recompile_limit，預設為 8；屆時我們將回退到 eager 模式。請注意，編譯快取是針對程式碼物件的，而不是幀；如果您動態建立函式的多個副本，它們將共享相同的程式碼快取。

引數

  • model (Callable 或 None) – 要最佳化的模組/函式

  • fullgraph (bool) – 如果為 False（預設值），torch.compile 會嘗試在函式中發現可編譯區域並進行最佳化。如果為 True，則我們要求整個函式都可以捕獲到一個圖中。如果這是不可能的（即存在圖中斷），則會引發錯誤。

  • dynamic (bool 或 None) – 使用動態形狀跟蹤。當此引數為 True 時，我們將預先嚐試生成一個儘可能動態的核心，以避免在尺寸改變時重新編譯。這可能並不總是有效，因為某些操作/最佳化會強制進行特化；使用 TORCH_LOGS=dynamic 來除錯過度特化。當此引數為 False 時，我們將永遠不會生成動態核心，我們將始終進行特化。預設值（None）是，我們自動檢測是否發生了動態性，並在重新編譯時編譯一個更動態的核心。

  • backend (str 或 Callable) –

    要使用的後端

  • mode (str) –

    可以是“default”、“reduce-overhead”、“max-autotune”或“max-autotune-no-cudagraphs”

    • “default”是預設模式，在效能和開銷之間取得了良好的平衡

    • “reduce-overhead”是一種透過 CUDA 圖降低 Python 開銷的模式，對於小批次很有用。降低開銷可能會以更多的記憶體使用為代價，因為我們將快取呼叫所需的記憶體工作區，以便在後續執行中不必重新分配它。降低開銷不一定有效；今天，我們只為不修改輸入的純 CUDA 圖降低開銷。在某些情況下 CUDA 圖也可能不適用；使用 TORCH_LOG=perf_hints 進行除錯。

    • “max-autotune”是一種利用 Triton 或模板進行矩陣乘法（在支援的裝置上）以及 Triton 進行卷積（在 GPU 上）的模式。它在 GPU 上預設啟用 CUDA 圖。

    • “max-autotune-no-cudagraphs”是一種與“max-autotune”類似的模式，但不使用 CUDA 圖。

    • 要檢視每種模式設定的確切配置，您可以呼叫 torch._inductor.list_mode_options()

  • options (dict) –

    傳遞給後端的選項字典。一些值得嘗試的亮點包括：

    • epilogue_fusion，它將逐點操作融合到模板中。需要同時設定 max_autotune

    • max_autotune，它將進行效能分析以選擇最佳的矩陣乘法配置

    • fallback_random，這在除錯精度問題時很有用

    • shape_padding，它會對矩陣形狀進行填充，以更好地在 GPU 上對齊載入，特別是對於張量核心

    • triton.cudagraphs，它將透過 CUDA 圖降低 Python 的開銷

    • trace.enabled，這是最有用的除錯標誌

    • trace.graph_diagram，它將顯示融合後圖的圖片

    • guard_filter_fn，用於控制將哪些 dynamo guard 與編譯一起儲存。這是一個不安全的特性，並且對於作為資料型別的 guard 沒有向後相容性保證。有關穩定的輔助函式，請參閱 torch.compiler 中的文件，例如： - torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe - torch.compiler.skip_guard_on_all_nn_modules_unsafe - torch.compiler.keep_tensor_guards_unsafe

    • 對於 inductor，您可以透過呼叫 torch._inductor.list_options() 檢視其支援的完整配置列表

  • disable (bool) – 將 torch.compile() 變成一個無操作，用於測試

示例

@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
def foo(x):
    return torch.sin(x) + torch.cos(x)
在此頁面上
顯示原始碼
PyTorch 庫