torch.compiler.set_stance

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

设置编译器的当前姿态。可以用作函数、上下文管理器或装饰器。不要在 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”: 在需要重新编译时以eager模式运行代码。如果输入有缓存的编译代码有效，仍然会使用它。

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

  • ”eager_then_compile”: 在第一次调用时以eager模式运行，然后在后续调用时进行编译。这对于动态形状是有益的，因为它允许从前两次调用中推断出动态性，而不是在第一次调用时浪费静态编译。

  • ”aot_eager_then_compile”: 使用 AOT eager 运行第一次调用，以从激活检查点获得内存优势，然后在后续调用时进行编译。与 eager_then_compile 类似，这通过避免初始静态编译来改善对动态形状的处理。

• skip_guard_eval_unsafe (bool) –

  一个标志，仅运行区分守卫。警告 - 此标志是不安全的，只有在您的设置满足以下条件时才应使用。

  torch.compile 使用守卫系统来支持重新编译并在运行时选择要运行的编译工件。这些守卫虽然高效，但会增加一些开销，这可能会影响在需要优化最小守卫处理时间的情况下性能。此 API 允许您禁用守卫评估，假设您已经使用足够多的输入“预热”了编译后的模型。该假设意味着，在预热阶段之后，不再需要进一步的重新编译。如果此假设失败，则有静默产生不正确结果的风险（因此 API 名称中包含“unsafe”一词）。

• force_backend (str | Callable[[...], Any] | None) – 如果 stance 为 “default”，则可以使用此参数强制 torch.compile 使用特定的后端。否则，会引发错误。