使用 Sublinear 进行显存优化

警告

Sublinear 仅支持在静态图模式下开启,参考 将动态图转为静态图(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):
     ...