评价此页

torch.compiler.set_stance#

torch.compiler.set_stance(stance='default', *, skip_guard_eval_unsafe=False, force_backend=None)[source]#

设置编译器的当前状态。可用作函数、上下文管理器或装饰器。请勿在 torch.compile 区域内使用此函数,否则将引发错误。

@torch.compile
def foo(x):
    ...

@torch.compiler.set_stance("force_eager")
def bar():
    # will not be compiled
    foo(...)

bar()

with torch.compiler.set_stance("force_eager"):
    # will also not be compiled
    foo(...)

torch.compiler.set_stance("force_eager")
# will also not be compiled
foo(...)
torch.compiler.set_stance("default")

# will be compiled
foo(...)
参数
  • stance (str) –

    要设置的编译器状态。有效值为:

    • “default”:默认状态,用于正常编译。

    • “force_eager”:忽略所有 torch.compile 指令。

    • “eager_on_recompile”:当需要重新编译时,以惰性模式运行代码。如果存在有效的已编译代码缓存,仍将使用它。

    • “fail_on_recompile”:在重新编译函数时引发错误。

    • “eager_then_compile”:首次调用以惰性模式运行,后续调用时编译。这有利于动态形状,因为它允许从前两次调用中推断动态性,而不是将第一次调用编译为静态。

    • “aot_eager_then_compile”:首次调用时使用 AOT 惰性模式以获得激活检查点的内存优势,后续调用时编译。与 eager_then_compile 类似,这可以通过避免初始静态编译来改善动态形状的处理。

  • skip_guard_eval_unsafe

    一个标志,用于仅运行区分性 guards。注意 - 此标志不安全,仅当您的设置满足以下条件时才应使用。

    torch.compile 使用 guards 系统来支持重新编译并在运行时选择要运行的编译工件。这些 guards 虽然高效,但会增加一些开销,在需要优化以最小化 guard 处理时间的场景中可能会影响性能。此 API 使您能够禁用 guard 评估,假设您已经用足够多样的输入预热了编译后的模型。此假设意味着,在预热阶段之后,将不再需要进一步的重新编译。如果此假设失败,则存在产生错误结果的风险(因此 API 名称中使用了“不安全”一词)。

  • force_backend – 如果 stance 为“default”,则此参数可用于强制 torch.compile 使用特定的后端。否则,将引发错误。