Memory 'leak' issue with MapTable and clearState

[achalddave]

Calling clearState() on MapTable seems to lead to a significant increase in memory usage for future iterations. I was unable to find the source, but the following script demonstrates the bug. (Of course, the exact memory amounts may not match across systems, but the relative amounts should be correct.)

--[[ Displays memory 'leak' with nn.MapTable after clearState() ]]--

local nn = require 'nn'
local cunn = require 'cunn'

model = nn.MapTable():add(nn.SpatialConvolution(3, 256, 3, 3, 1, 1, 1, 1, 1))
                     :cuda()
i = {torch.rand(30, 3, 224, 224):cuda(), torch.rand(30, 3, 224, 224):cuda()}

function check_mem() os.execute('nvidia-smi | grep luajit') end
-- Train two iterations without clear state:
print('Before training 1'); check_mem() -- 277 MiB
o = model:forward(i)
print('After forward 1'); check_mem() -- 3254 MiB

-- Train another iteration:
print('Before training 2'); check_mem() -- 3254 MiB
o = model:forward(i)
print('After forward 2'); check_mem() -- 3254 MiB

-- Clear state:
model:clearState()
collectgarbage()
collectgarbage()

-- Train a final iteration before clearState. This final forward call causes an
-- increase in memory usage for the rest of the program!
print('Before training 3 (after clearState)'); check_mem() -- 3254 MiB
o = model:forward(i)
print('After forward 3 (after clearState)'); check_mem() -- 4724 MiB!

-- Garbage collection doesn't fix it.
collectgarbage()
collectgarbage()
print('After collectgarbage()'); check_mem() -- 4724 MiB!

The script is at: https://gist.github.com/achalddave/6ac8390e06a23ecc6d67e3fa22ef0f04

A few notes:

• The issue does not seem to occur if I attempt to only forward an input table with only one element.
• The issue does not occur if nn.MapTable is removed (and replaced with just a single SpatialConvolution operating on a single input tensor).

[achalddave]

Update: This issue is fixed if the modules in MapTable have clearState called before they are removed from MapTable, here:

nn/MapTable.lua

Line 80 in acc6b8c

function MapTable:clearState()