使用 Sublinear 进行显存优化

Warning

Sublinear 仅支持在静态图模式下开启，参考 Convert dynamic graphs to static graphs (Trace) 。

MegEngine 通过引入 Sublinear [1] 技术来进行静态图下的显存优化。

[1]

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.

用户在编译静态图时使用 SublinearMemoryConfig 设置 trace 的参数 sublinear_memory_config, 就可以打开 Sublinear 优化：

from megengine.jit import trace, SublinearMemoryConfig

config = SublinearMemoryConfig()

@trace(symbolic=True, sublinear_memory_config=config)
def train_func(data, label, * , net, optimizer, gm):
     ...