High memory usage while training

NgPDat: <p>Hello all,<br>
I train a simple RNN network to predict a label on each input timestep on a huge random dataset.<br>
I record memory usage while training, and notice that it is increasing linearly with dataset size:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/1X/61f5e7b59b3dfba741c7999d644ce9d536db3174.png" data-download-href="https://discuss.pytorch.org/uploads/default/61f5e7b59b3dfba741c7999d644ce9d536db3174" title="mem-graph.png"><img src="https://discuss.pytorch.org/uploads/default/original/1X/61f5e7b59b3dfba741c7999d644ce9d536db3174.png" width="666" height="500" data-dominant-color="FDFDFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">mem-graph.png</span><span class="informations">800×600 3.74 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>(VSIZE = Virtual Memory recorded by Ubuntu, %MEM: How much % RAM it takes, x-axis = time in second)<br>
My training script for reference:</p>
<pre><code class="lang-auto">class testNet(nn.Module):
    def __init__(self):
        super(testNet, self).__init__()
        self.rnn = nn.RNN(input_size=200, hidden_size=1000, num_layers=1)
        self.linear = nn.Linear(1000, 100)

    def forward(self, x, init):
        x = self.rnn(x, init)[0]
        y = self.linear(x.view(x.size(0)*x.size(1), x.size(2)))
        return y.view(x.size(0), x.size(1), y.size(1))

net = testNet()
init = Variable(torch.zeros(1, 4, 1000))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

total_loss = 0.0
for i in range(10000): #10000 mini-batch
    input = Variable(torch.randn(1000, 4, 200)) #Seqlen = 1000, batch_size = 4, feature = 200
    target = Variable(torch.LongTensor(4, 1000).zero_())

    optimizer.zero_grad()
    output = net(input, init)
    loss = criterion(output.view(-1, output.size(2)), target.view(-1))
    loss.backward()
    optimizer.step()
    total_loss += loss[0]

print(total_loss)
</code></pre>
<p>I expect memory usage not increasing per mini-batch. What might be the problem? (Correct me if my script is wrong)</p>

smth: <p>hi <a class="mention" href="/u/ngpdat">@NgPDat</a></p>
<p>I’m trying to reproduce your results.</p>
<p>Can you tell me the units of VSIZE? Is it bytes?<br>
And %MEM, is it a percentage of the system memory?</p>
<p>So far, my run is pretty stable at around 105MB, after 400 mini-batches, I will wait for some time.</p>

apaszke: <p>I think I see the problem. You have to remember that <code>loss</code> is a Variable, and indexing Variables, always returns a Variable, even if they’re 1D! So when you do <code>total_loss += loss[0]</code> you’re actually making <code>total_loss</code> a Variable, and adding more and more subgraphs to its history, making it impossible to free them, because you’re still holding a reference. Just replace <code>total_loss += loss[0]</code> with <code>total_loss += loss.data[0]</code> and it should be back to normal.</p>

NgPDat: <p>Work like a charm! Thank you.<br>
I think I have better understanding of Variable now.</p>

NgPDat: <p>VSIZE is in kilobytes.<br>
Yes, %MEM is percentage of the system memory.</p>
<p>The whole script for recording memory usage is from here: <a href="http://stackoverflow.com/questions/7998302/graphing-a-processs-memory-usage" rel="nofollow noopener">http://stackoverflow.com/questions/7998302/graphing-a-processs-memory-usage</a></p>

ado_sar: <p>Why we add subgraphs to its history? Is it because <code>loss</code> still requires grad after <code>loss.backward()</code>?</p>