Python PyTorch Proxy用法及代碼示例

本文簡要介紹python語言中 torch.fx.Proxy 的用法。

用法:

class torch.fx.Proxy(node, tracer=None)

Proxy 對象是 Node 包裝器，在符號跟蹤期間流經程序，並記錄它們接觸到不斷增長的 FX Graph 的所有操作(torch 函數調用、方法調用、運算符)。

如果您正在執行圖形轉換，您可以將自己的 Proxy 方法包裝在原始 Node 周圍，以便您可以使用重載的運算符向 Graph 添加其他內容。

Proxy 對象無法迭代。換句話說，如果在循環中使用 Proxy 或作為 *args /**kwargs 函數參數，符號跟蹤器將引發錯誤。

解決這個問題的主要方法有兩種： 1. 將無法追蹤的邏輯分解為頂級函數並在其上使用fx.wrap。 2. 如果控製流是靜態的(即循環行程計數基於某個超參數)，代碼可以保持在其原始位置並重構為：

for i in range(self.some_hyperparameter):
    indexed_item = proxied_value[i]

有關代理內部的更詳細說明，請查看 torch/fx/OVERVIEW.md 中的 “Proxy” 部分

注意

保證此 API 的向後兼容性。