使用 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):
...