functorch.compile.memory_efficient_fusion¶

functorch.compile.memory_efficient_fusion(fn, static_argnums=None, **kwargs)[source]¶

包装函数覆盖了aot_function()和aot_module()，以执行内存高效融合。它使用min_cut_rematerialization_partition()分区器来执行高效的重计算。它使用NVFuser来编译生成的前向和后向图。

警告

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

Parameters

fn (Union[Callable, nn.Module]) – 一个Python函数或nn.Module，它接受一个或多个参数。必须返回一个或多个张量。
static_argnums (Optional[Tuple[Int]]) – 一个可选的整数元组，用于将函数的参数标记为静态。
**kwargs – 您想要对设置进行的任何其他覆盖

Returns

返回一个Callable或nn.Module，它保留了原始fn的急切行为，但其前向和后向图已经通过重新计算优化，并且这些图已经使用nvfuser进行了编译。