Skip to content

Sublinear memory optimization for deep learning, reduce GPU memory cost to train deeper nets

License

Notifications You must be signed in to change notification settings

dmlc/mxnet-memonger

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

9 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

MXNet Memory Monger

This project contains a 150 lines of python script to give sublinear memory plans of deep neural networks. This allows you to trade computation for memory and get sublinear memory cost, so you can train bigger/deeper nets with limited resources.

Reference Paper

Training Deep Nets with Sublinear Memory Cost Arxiv 1604.06174

How to Use

This code is based on MXNet, a lightweight, flexible and efficient framework for deep learning.

  • Configure your network as you normally will do using symbolic API
  • Give hint to the allocator about the possible places that we need to bookkeep computations.
    • Set attribute mirror_stage='True', see example_resnet.py
    • The memonger will try to find possible dividing points on the nodes that are annotated as mirror_stage.
  • Call memonger.search_plan to get an symbolic graph with memory plan.
import mxnet as mx
import memonger

# configure your network
net = my_symbol()

# call memory optimizer to search possible memory plan.
net_planned = memonger.search_plan(net)

# use as normal
model = mx.FeedForward(net_planned, ...)
model.fit(...)

Write your Own Memory Optimizer

MXNet's symbolic graph support attribute to give hint on whether (mirror attribute) a result can be recomputed or not. You can choose to re-compute instead of remembering a result for less memory consumption. To set output of a symbol to be re-computable, use

sym._set_attr(force_mirroring='True')

mxnet-memonger actually use the same way to do memory planning. You can simply write your own memory allocator by setting the force_mirroring attribute in a smart way.

About

Sublinear memory optimization for deep learning, reduce GPU memory cost to train deeper nets

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages