Accumulating model output blows up cuda memory?

Goldname: <p>As I accumulate the model outputs by concatenating them, my CUDA memory grows significantly more than the size of my torch tensors: “out” and “lab”.<br>
I am accumulating the output because I need gradient accumulation downstream.</p>
<p>I’ve simplified my code down to the pytorch classifier example code:</p>
<pre><code class="lang-auto">import torch
import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

batch_size = 16

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


device = "cuda" if torch.cuda.is_available() else "cpu"
out = torch.empty(0, 10).to(device)
lab = torch.empty(0, dtype=torch.int64).to(device)
net = Net().to(device)
for epoch in range(2):  # loop over the dataset multiple times
    for i, (inputs, labels) in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = net(inputs)
        out = torch.cat((out, outputs), dim=0)
        lab = torch.cat((lab, labels), dim=0)

</code></pre>
<p>These tensors should be able to easily fit in memory but CUDA is quickly running out of memory.</p>
<p>In the sample code I gave, everything fits into memory because it’s a toy example, but the CUDA memory usage still grows significantly. After the code finishes nvidia-smi shows that my process is using 5868MiB. Meanwhile “lab” and “out” are float32 with size 100000 x 10, which should only be a few MiB.</p>
<p>I am guessing this might have something to do with the computational graph, but if I instead use the standard approach to gradient accumulation where the gradients are accumulated instead, this error doesn’t occur. Why is that?</p>

ptrblck: <aside class="quote no-group" data-username="Goldname" data-post="1" data-topic="197063">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/db5fbb/48.png" class="avatar"> Goldname:</div>
<blockquote>
<p>As I accumulate the model outputs by concatenating them, my CUDA memory grows significantly more than the size of my torch tensors: “out” and “lab”.</p>
</blockquote>
</aside>
<p>This is expected since you are storing the entire computation graphs with these tensors including the intermediate activations needed to compute the gradients during the backward pass.</p>
<aside class="quote no-group" data-username="Goldname" data-post="1" data-topic="197063">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/db5fbb/48.png" class="avatar"> Goldname:</div>
<blockquote>
<p>I am accumulating the output because I need gradient accumulation downstream.</p>
</blockquote>
</aside>
<p>You can call <code>.backward()</code> on each loss and accumulate the gradients directly in the <code>.grad</code> attribute. If you don’t want to do it for some reason and want to call <code>backward</code> only once, your memory increase is expected as described above.</p>

Goldname: <p>If I instead scaled up the batch size in the data loader, would this problem also happen? I thought the memory increase with a large batch size was solely due to the batch itself and not an increase in the computational graph.</p>

ptrblck: <p>Yes, you should see a similar memory increase.</p>
<aside class="quote no-group" data-username="Goldname" data-post="3" data-topic="197063">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/db5fbb/48.png" class="avatar"> Goldname:</div>
<blockquote>
<p>memory increase with a large batch size was solely due to the batch itself and not an increase in the computational graph.</p>
</blockquote>
</aside>
<p>The computation graph stores forward activations, which are needed for the gradient computation, and thus increases the memory. The actual nodes and metadata should be stored on the host and should not increase device memory.</p>

Goldname: <p>Thanks a lot! So does the computation graph get duplicated in size if I duplicate the number of items in the batch?</p>

ptrblck: <p>Yes, since all intermediates will use the same <code>batch_size</code> in their shape and will thus see the same factor in their size increase.</p>