functorch.compile.memory_efficient_fusion

functorch.compile.memory_efficient_fusion(fn, static_argnums=None, **kwargs)

对 aot_function() 和 aot_module() 的包装函数，用于执行内存高效融合。它使用 min_cut_rematerialization_partition() 分区器来执行高效的重新计算。它使用 NVFuser 来编译生成的正向和反向图。

警告

此 API 为实验性 API，可能会有所改变。

参数

• fn (Union[Callable, nn.Module]) – 一个 Python 函数或一个 nn.Module，它接受一个或多个参数。必须返回一个或多个张量。

• static_argnums (Optional[Tuple[Int]]) – 用于标记函数参数为静态的可选元组。

• **kwargs – 您希望对设置进行的任何其他覆盖。

返回值

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