Calling loss.backward() reduce memory usage?

gwding: <p>It seems that calling loss.backward() help save memory, which is not very intuitive to me.</p>
<pre><code>import sys
import torch
import torch.utils
import torch.nn as nn
from torch.autograd import Variable
from torchvision import models

try:
    import gpustat
except ImportError:
    raise ImportError("pip install gpustat")


def show_memusage(device=0):
    gpu_stats = gpustat.GPUStatCollection.new_query()
    item = gpu_stats.jsonify()["gpus"][device]
    print("{}/{}".format(item["memory.used"], item["memory.total"]))


device = 0
show_memusage(device=device)

torch.cuda.set_device(device)
model = models.resnet101(pretrained=False)
model.cuda()
criterion = nn.CrossEntropyLoss()

volatile = False

show_memusage(device=device)
for ii in range(3):
    inputs = torch.randn(20, 3, 224, 224)
    labels = torch.LongTensor(range(20))
    inputs = inputs.cuda()
    labels = labels.cuda()
    inputs = Variable(inputs, volatile=volatile)
    labels = Variable(labels, volatile=volatile)

    print "before run model:",
    show_memusage(device=device)
    outputs = model(inputs)
    print "after run model:",
    show_memusage(device=device)

    if bool(int(sys.argv[1])):
        loss = criterion(outputs, labels)
        print "before backward:",
        show_memusage(device=device)
        loss.backward()
        print "after backward:",
        show_memusage(device=device)
    print
</code></pre>
<p>running “python test_dsetloadermem.py 0” gives</p>
<pre><code>mem used / mem total
612/8110
1045/8110
before run model: 1055/8110
after run model: 3699/8110

before run model: 3711/8110
after run model: 6215/8110

before run model: 6215/8110
after run model: 6215/8110
</code></pre>
<p>running “python test_dsetloadermem.py 1” gives<br>
mem used / mem total<br>
618/8110<br>
1051/8110<br>
before run model: 1063/8110<br>
after run model: 3699/8110<br>
before backward: 3699/8110<br>
after backward: 4131/8110</p>
<pre><code>before run model: 4131/8110
after run model: 4131/8110
before backward: 4131/8110
after backward: 4131/8110

before run model: 4131/8110
after run model: 4131/8110
before backward: 4131/8110
after backward: 4131/8110
</code></pre>
<p>I’ve also noticed that</p>
<pre><code>del outputs
del inputs
del labels
</code></pre>
<p>can help avoid unnecessary memory usage.</p>
<p>Could you give any insight on why loss.backward() does something similar?<br>
Thanks!</p>

apaszke: <p>It’s because of the scoping rules in Python. When you do this:</p>
<pre><code class="lang-python">while True:
    loss = model(output)
</code></pre>
<p>it will always use 2x the memory that is needed to compute the model, because the reference to the <code>loss</code> form the previous iteration won’t be overwritten (and thus the graph with all the buffers it holds won’t be freed) until this iteration completes. So you’ll effectively end up holding to two graphs. This is why you should use <code>volatile=True</code> inputs when only doing inference. Once you add <code>.backward()</code> the buffers will be freed in the process of computing the derivatives, limiting the memory usage to that of a single graph (the old graph will be still kept around, but it won’t be holding to any memory).</p>
<p>This is also why <code>del</code> helps reduce memory usage. This loop will only keep at most a single graph alive, because the <code>loss</code> is created and disposed within a single iteration.</p>
<pre><code class="lang-python">while True:
    loss = model(input)
    del loss # This frees the graph
</code></pre>

gwding: <p>Thanks for the explanation! that helps a lot!</p>

ptrblck: <p>The forward pass will create the computation graph, which is used in the backward pass to calculate the gradients.<br>
<code>loss</code> is attached to this graph and the graph won’t be freed until <code>loss</code> is deleted of goes out of scope.<br>
Of course the computation graph holds references to intermediate tensors, which are also needed to compute the gradients, and will thus use memory.</p>
<p>Given this simple loop:</p>
<pre><code class="lang-python">while True:
    loss = model(output)
</code></pre>
<p>These steps will happen:</p>
<ul>
<li>
<code>model</code> and <code>output</code> are created and thus will use memory</li>
<li>
<code>loss = model(output)</code> represents the first forward pass. The computation graph as well as the intermediate tensors will be created. Memory usage increases (model + output + loss0 + intermediates0)</li>
<li>the iteration is done and the next one will start. Current memory usage is still (model + output + loss0 + intermediates0)</li>
<li>the next iteration will start and another forward call will be kicked off. Memory usage during the forward pass and before the assignment of <code>loss</code>: (model + output + intermediates1 + loss0 + intermediates0) Note that we end up creating two graphs at this point.</li>
<li>the new <code>loss</code> will be assigned to the variable <code>loss</code>. The old <code>loss</code> (loss0) will be freed and thus also intermediates0 and the corresponding graph</li>
</ul>