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 在生成的图中将排除哪些函数，并强制在这些函数处断开图。

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=None, guard_fail_fn=None, disable=False, dynamic=False)

TorchDynamo 的主要入口点。执行图捕获并调用 backend() 以优化提取的图。

Parameters:

• backend – 两件事情之一： - 要么，一个函数/可调用对象，接受一个 torch.fx.GraphModule 和 example_inputs，并返回一个运行图更快的 Python 可调用对象。 还可以通过设置 backend_ctx_ctor 属性，为后端提供额外的上下文，例如 torch.jit.fuser("fuser2")。 有关用法，请参见 AOTAutogradMemoryEfficientFusionWithContext。 - 或者，torch._dynamo.list_backends() 中的一个字符串后端名称

• nopython – 如果为 True，图断点将被视为错误，并且将存在一个完整的程序图。

• 禁用 – 如果为 True，将此装饰器转换为无操作

• 动态 – 如果为 True，启用动态形状支持

示例用法：

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

torch._dynamo.optimize_assert(backend, *, hooks=Hooks(guard_export_fn=None, guard_fail_fn=None), 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=None)

跳过与函数代码相关的帧，但仍处理递归调用的帧

class torch._dynamo.OptimizedModule(mod, dynamo_ctx)

包装原始的 nn.Module 对象，并随后修补其 forward 方法以使用优化后的 self.forward 方法。

torch._dynamo.register_backend(compiler_fn=None, name=None, tags=())

用于将给定编译器添加到注册表中的装饰器，以允许使用字符串简写调用 torch.compile。注意：对于未默认导入的项目，直接传递一个函数作为后端可能更简单，而不需要使用字符串。

Parameters:

• compiler_fn (可选[Callable[[GraphModule, List[Tensor]], CompiledFn]]) – 接受FX图和假张量输入的可调用对象

• 名称 (可选[字符串]) – 可选名称，默认为 compiler_fn.__name__

• 标签 (序列[字符串]) – 可选的字符串标签集，用于对后端进行分类

torch._dynamo.list_backends(exclude_tags=('debug', 'experimental'))

返回可以传递给的合法字符串：

torch.compile(…, backend=”name”)