functorch.compile.aot_module

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

使用torch dispatch跟踪机制跟踪mod的前向和后向图。它是一个包装函数，底层使用aot_function()来执行跟踪和编译。

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

警告

此API是实验性的，可能会发生变化。

Parameters

• mod (可调用的) – 一个 nn.Module 模块。

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

• kwargs – 要传递给 aot_function() 的 kwargs

Returns

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