显著减少MXNet显存占用--MXNet Memory Monger

最近在训练deepinsight/insightface 的时候遇到一个问题,官方提供的设置是batchsize 128, 而我的机器内存不够,只能设为64来训练,但是batch_size不一样对最终训练的结果是有一部分影响的。为了达到最佳效果,就要解决out of memory的问题,怎么办呢?于是发现了这个:dmlc/mxnet-memonger

使用方法

使用方法非常简单,下载该github项目中的 memonger.py, 将其放到训练代码同文件夹下,然后在训练代码中添加memonger.search_plan:

import memonger

net = my_symbol() net_planned = memonger.search_plan(net)  
# 之后的部分和之前一样就可以了 
model = mx.FeedForward(net_planned,  ...) 
model.fit(...)

效果

我在insightface中加了memonger后,之前训练64的batchsize需要10G内存,现在训练128的batch也只要8G内存了,而且,128的batch_size效果的确比64好..

推荐阅读更多精彩内容