OptNet：面向torch的神经网络内存优化工具

OptNet - reducing memory usage in torch neural networks

Memory optimizations for torch neural networks.

Heavily inspired from the Optimizer from https://github.com/facebook/fb-caffe-exts

How does it work ?

It goes over the network and verify which buffers can be reused. Currently, it only supports evaluation mode, but training mode will soon be included.

Here is a list of currently tested modules. Numbers are for CPU version, with batch size of 1, in the format (total memory used, memory used for the outputs):

Note that most of the used memory goes to the convolution buffers from nn .

In a more realistic setup where we use cudnn and batch size of 128, the gains are way more significant. The total memory usage is shown in the following table:

Visualizing the memory reuse

We can analyse the sharing of the internal buffers by looking at the computation graph of the network before and after the sharing.

For that, we have the graphgen(net, input, opts) function, which creates the graph corresponding to the network net . The generated graph contains the storage id of each output , and same colors means same storage.

Note that net is a nn model, and not a nngraph network. This allows us to use optnet.graphgen to generate graph visualizations of nn networks without having to use nngraph .

Let's have a look:

-- some handy models are defined in optnet.models -- like alexnet, googlenet, vgg and resnet models = require 'optnet.models' modelname = 'googlenet' net, input = models[modelname]()  generateGraph = require 'optnet.graphgen'  g = generateGraph(net, input)  graph.dot(g,modelname,modelname)

This generates the following graph:

Now what happens after we optimize the network ? Check the colors and the storage ids.

models = require 'optnet.models' modelname = 'googlenet' net, input = models[modelname]()  opts = {inplace=true, reuseBuffers=true}  generateGraph = require 'optnet.graphgen'  optnet = require 'optnet'  optnet.optimizeMemory(net, input, opts)  g = generateGraph(net, input)  graph.dot(g,modelname..'_optimized',modelname..'_optimized')

Counting the amount of saved memory

We can also provide a function to compute the amount of memory used by the network in bytes, which allows us to check the amount of saved memory. To count the only the memory used by the output state variables of each module, pass the option {countBuffers=false} .

Here is an example

optnet = require 'optnet'  models = require 'optnet.models' modelname = 'googlenet' net, input = models[modelname]()  -- count the memory used by all the buffers opts = {countBuffers=true}  mem1 = optnet.countUsedMemory(net, input, opts)  optnet.optimizeMemory(net, input)  mem2 = optnet.countUsedMemory(net, input, opts)  optnet.removeOptimization(net)  mem3 = optnet.countUsedMemory(net, input, opts)  print('Before optimization        : '.. mem1/1024/1024 .. ' MBytes') print('After optimization         : '.. mem2/1024/1024 .. ' MBytes') print('After removing optimization: '.. mem3/1024/1024 .. ' MBytes')