torch._dynamo

警告

此模块是早期原型，可能会发生更改。

torch._dynamo。allow_in_graph（fn）

自定义 TorchDynamo 将包含在生成的 图。类似于 torch.fx.wrap（）。

torch._dynamo.allow_in_graph(my_custom_function)

@torch._dynamo.optimize(...)
def fn(a):
    x = torch.add(x, 1)
    x = my_custom_function(x)
    x = torch.add(x, 1)
    return x

fn(...)

将捕获包含 my_custom_function（） 的单个图形。

torch._dynamo。disallow_in_graph（fn）

自定义 TorchDynamo 将在生成的 graph 并强制 graph break on。

torch._dynamo.disallow_in_graph(torch.sub)

@torch._dynamo.optimize(...)
def fn(a):
    x = torch.add(x, 1)
    x = torch.sub(x, 1)
    x = torch.add(x, 1)
    return x

fn(...)

将破坏 torch.sub 上的图形，并给出两个图形，每个图形都有一个 单个 torch.add（） 操作。

torch._dynamo。graph_break（

强制图形中断

torch._dynamo。optimize（backend='inductor'， *， nopython=False， guard_export_fn=无， guard_fail_fn=无， disable=False， dynamic=False）

TorchDynamo 的主入口点。执行图形捕获并调用 backend（） 来优化提取的图形。

参数

• backend – 以下两项之一： - 一个函数/可调用对象接受 torch.fx.GraphModule 和 example_inputs并返回一个运行 graph 更快的速度。 还可以为后端提供额外的上下文，例如 torch.jit.fuser（“fuser2”） 的 backend_ctx_ctor 属性。 有关用法，请参阅 AOTAutogradMemoryEfficientFusionWithContext 。 - 或者，在 torch._dynamo.list_backends（） 中输入字符串后端名称

• nopython – 如果为 True，则图形中断将是错误，并且会 是单个全程序图。

• disable – 如果为 True，则将此装饰器转换为 no-op

• dynamic – 如果为 True，则打开动态形状支持

示例用法：

@torch._dynamo.optimize()
def toy_example(a, b):
    ...

torch._dynamo。optimize_assert（backend， *， hooks=Hooks（guard_export_fn=无， guard_fail_fn=无）， export=False， dynamic=False）

同 torch._dynamo.optimize（backend， nopython=True）

torch._dynamo。run（fn=None）

不做任何动态编译，只使用之前的优化

torch._dynamo。disable（fn=None）

用于禁用 TorchDynamo 的装饰器和上下文管理器

torch._dynamo。reset（

清除所有编译缓存并恢复初始状态

torch._dynamo。skip（fn=无）

跳过与函数代码关联的帧，但仍以递归方式处理 调用的帧

torch._dynamo 类。OptimizedModule（mod， dynamo_ctx）

将原始 nn.Module 对象，并在 forward 方法转换为优化的 self.forward 方法。

torch._dynamo。register_backend（compiler_fn=无，name=无，tags=（）））

Decorator 将给定的编译器添加到注册表中，以允许使用字符串速记调用 torch.compile 。注意：对于项目 imported 时，直接传递函数可能更容易 作为后端，而不是使用字符串。

参数

• compiler_fn （Optional[Callable[[GraphModule， List[Tensor]]， CompiledFn]]） – 采用 FX 图和假张量输入的可调用

• name （Optional[str]） – 可选名称，默认为 compiler_fn.__name__

• tags （Sequence[str]） – 用于对后端进行分类的可选字符串标签集

torch._dynamo。list_backends（exclude_tags=（'debug'， 'experimental'））

返回可传递给的有效字符串：

torch.compile（...， backend=“名称”）