torch.compiler.allow_in_graph¶
- torch.compiler.allow_in_graph(fn)[源代码]¶
告诉编译器前端 (Dynamo) 跳过对函数的符号内省,而是在遇到时直接将其写入图形。
如果您正在使用
torch.compile()
(使用后端“inductor”(默认值)),或者torch.export.export()
,并且尝试在整个跟踪中黑盒化 Python 函数,请勿使用此 API。相反,请创建一个自定义运算符(请参阅 PyTorch 自定义运算符着陆页)。警告
如果您是典型的 torch.compile 用户(例如,您将 torch.compile 应用于模型以使其运行更快),您可能不希望使用此函数。
allow_in_graph()
是一把双刃剑,因为它会跳过负责执行安全检查(图形中断、处理闭包等)的编译器前端 (Dynamo)。不正确的使用会导致难以调试的静默错误问题。给定一个没有 allow_in_graph 装饰器的 Python 函数,torch.compile 的常规执行会跟踪该函数。
allow_in_graph()
会更改它,以便前端不会跟踪函数内部,但编译器后端仍然会跟踪它。将其与自定义运算符进行比较,自定义运算符将整个 torch.compile 堆栈中的函数视为黑盒。下表比较了这些机制。机制
前端 (Dynamo)
后端 (AOTAutograd+Inductor)
无装饰器
跟踪内部
跟踪内部
allow_in_graph
不透明的可调用对象
跟踪内部
自定义运算符
不透明的可调用对象
不透明的可调用对象
allow_in_graph()
的一个常见用例是作为编译器前端的逃生舱口:如果您知道函数相对于编译堆栈的下游组件(AOTAutograd 和 Inductor)工作正常,但 Dynamo 中存在阻止其正确符号内省函数的错误(或者您的代码是用 C/C++ 编写的,因此无法使用 Dynamo 进行内省),那么您可以使用allow_in_graph()
来装饰该函数以绕过 Dynamo。我们要求
fn
遵守以下限制。未遵守会导致未定义的行为fn
的输入必须是 FX 图表中的可代理类型。有效类型包括:张量/int/bool/float/None/List[张量?]/List[int?]/List[float?] Tuple[张量?,…]/Tuple[int?,…]/Tuple[float?,…]/torch.dtype/torch.devicefn
的输出必须是 FX 图表中的可代理类型(请参阅上一条要点)。fn
内部使用的所有张量必须直接作为输入传递给fn
(而不是被捕获的变量)。
- 参数
fn – 表示要包含在图形中的函数的可调用对象。如果
fn
是一个可调用对象的列表或元组,它会递归地将allow_in_graph()
应用于每个函数,并返回一个包含修改后的函数的新列表或元组。
示例
torch.compiler.allow_in_graph(my_custom_function) @torch.compile(...) def fn(a): x = torch.add(x, 1) x = my_custom_function(x) x = torch.add(x, 1) return x fn(...)
将捕获包含
my_custom_function()
的单个图形。