functorch.compile.aot_module

functorch.compile.aot_module(mod, *args, **kwargs)

使用 torch 派发跟踪机制跟踪 mod 的前向和后向图。它是一个包装函数，在内部使用 aot_function() 执行跟踪和编译。

aot_module() 将 nn.Module 的参数和缓冲区提升为新可调用的输入，然后通过 aot_function() 进行编译。

警告

此 API 处于实验阶段，可能会发生变化。

参数

• mod (Callable) – 一个 nn.Module 模块。

• args – 传递给 aot_function() 的参数

• kwargs – 传递给 aot_function() 的关键字参数

返回值

返回一个 nn.Module，它保留了原始 mod 的急切行为，但前向和后向图已编译。