torchrec.fx

Torchrec 跟踪器

torchrec 的自定义 FX 跟踪器

请参阅 Torch.FX 文档

torchrec.fx.tracer

class 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，此函数将返回一个通过跟踪 root 中看到的操作构造的 GraphModule。

concrete_args 允许您部分专门化您的函数，无论它是为了移除控制流还是数据结构。

参数：:

• root (Union[torch.nn.Module, Callable]) – 要跟踪并转换为图形表示的模块或函数。

• concrete_args (Optional[Dict[str, any]]) – 要部分专门化的输入

返回：:

从 root 中记录的操作创建的模块。

返回类型：:

GraphModule

模块内容

Torchrec 跟踪器

torchrec 的自定义 FX 跟踪器

请参阅 Torch.FX 文档