torchrec.fx

Torchrec 追蹤器

用於 torchrec 的自訂 FX 追蹤器

請參閱 Torch.FX 文件

torchrec.fx.tracer

類別 torchrec.fx.tracer.Tracer(leaf_modules: Optional[List[str]] = None)

基底： Tracer

用於 torchrec 的自訂 FX 追蹤器

請參閱 Torch.FX 文件

我們建立了一個自訂 FX 追蹤器來追蹤基於 torchrec 的模型。自訂追蹤器會處理 Python 泛型類型（即 NoWait[T]、Awaitable[T]），並在需要時將其降低為 TorchScript

create_arg(a: Any) → Optional[Union[Tuple[Any, ...], List[Any], Dict[str, Any], slice, range, Node, str, int, float, bool, complex, dtype, Tensor, device, memory_format, layout, OpOverload, SymInt, SymBool, SymFloat]]

一種方法，用於指定在準備用作 Graph 中節點的參數值時，追蹤的行為。

除了預設追蹤器之外，還增加了對 NoWait 類型的支援

參數：:

a (Any) – 要在 Graph 中發出為 Argument 的值。

回傳：:

轉換為適當 Argument 的值 a

回傳類型：:

Argument

is_leaf_module(m: Module, module_qualified_name: str) → bool

覆寫 FX 定義以包含量化的嵌入包

path_of_module(mod: Module) → str

允許追蹤未註冊的模組。這通常用於讓表格批次嵌入看起來像 nn.EmbeddingBags

trace(root: Union[Module, Callable[[...], Any]], concrete_args: Optional[Dict[str, Any]] = None) → Graph

注意

保證此 API 的回溯相容性。

torchrec.fx.tracer.is_fx_tracing() → bool

torchrec.fx.tracer.symbolic_trace(root: Union[Module, Callable], concrete_args: Optional[Dict[str, Any]] = None, leaf_modules: Optional[List[str]] = None) → GraphModule

符號追蹤 API

給定一個 nn.Module 或函數實例 root，此函數將返回一個 GraphModule，該模組是通過記錄追蹤 root 時看到的操作來構建的。

concrete_args 允許您部分特化您的函數，無論是移除控制流程還是數據結構。

參數：:

• root (Union[torch.nn.Module, Callable]) – 要追蹤並轉換為圖形表示的模組或函數。

• concrete_args (Optional[Dict[str, any]]) – 要部分特化的輸入

回傳：:

從 root 記錄的操作創建的模組。

回傳類型：:

GraphModule

模組內容

Torchrec 追蹤器

用於 torchrec 的自訂 FX 追蹤器

請參閱 Torch.FX 文件