Interesting Issue with Automatic GPU Garbage Collection

username: Epoching
cooked: <p>Hello!</p>
<p>I recently found something interesting when trying to train a Mask RCNN. For those unfamiliar with this type of model, there are 2 inputs:</p>
<ol>
<li>A list of images (variable size)</li>
<li>A list of dictionaries (where each dictionary contains 3 tensors: bounding boxes, segmentation, and class label)</li>
</ol>
<p>More information regarding input data format can be found <a href="https://pytorch.org/docs/stable/torchvision/models.html#mask-r-cnn" rel="nofollow noopener">here in the documentation</a></p>
<p>The output is a dictionary of losses (classification loss, segmentation loss, etc). There’s a total of 5 keys/losses. I’m currently summing the losses into one tensor and then calling backward on the summed total loss. (The sum operation will be captured in the tensor graph).</p>
<p>Anyway, here is my problem that I ran into:</p>
<p>Below is some code for a simple training epoch:</p>
<pre><code class="lang-auto">for data in enumerate(someDataLoader):
    # get the inputs; data is a list of [inputs, labels]
    images, targets = data
        
    scenes = images = list(image.to(device) for image in images)
    for t in targets:
        t['boxes'] = t['boxes'].to(device)
        t['labels'] = t['labels'].to(device)
        t['masks'] = t['masks'].to(device)

    # forward + backward + optimize
    loss_dict = maskrcnn_model(scenes, targets)
    
    tot_loss = sum(loss for loss in loss_dict.values())
    
    # zero the parameter gradients
    maskrcnn_opt.zero_grad()
    tot_loss.backward()
    maskrcnn_opt.step()
</code></pre>
<p>When running this training loop, I notice that the GPU memory is increasing as more mini-batches are fed. (For example, at the start of training there would be a 3 GB memory usage, after 250 mini-batches, the GPU is at 12 GB memory usage, steadily increasing by 300 MB every 2 seconds). Batch size is fairly small.</p>
<p>To fix this, I looked at some online examples of training a Mask RCNN model. Everything was the same except this portion of the training loop:</p>
<pre><code class="lang-auto">for t in targets:
    t['boxes'] = t['boxes'].to(device)
    t['labels'] = t['labels'].to(device)
    t['masks'] = t['masks'].to(device)
</code></pre>
<p>Instead of doing it the way I did above, they used a list comprehension and basically made a new dictionary for each target, and saved into a separate list. An example of that would be either of the following code snippets below (I find the top code snippet to be more read-able, but they both do the same thing).</p>
<pre><code class="lang-auto">newTargs = []
for t in targets:
    x = {}
    x['boxes'] = t['boxes'].to(device)
    x['labels'] = t['labels'].to(device)
    x['masks'] = t['masks'].to(device)
    newTargs.append(x)
</code></pre>
<pre><code class="lang-auto">newTargs = [{k: v.to(device) for k, v in t.items()} for t in targets]
</code></pre>
<p>After making a new list of dictionaries, <code>newTargs</code> is then used in the forward pass of the model, rather than <code>targets</code>.</p>
<p>E.g. <code>loss_dict = maskrcnn_model(images, newTargs)</code>.</p>
<p>This small change in code maintains GPU memory of 3 GB throughout the training epoch.</p>
<p>Any thoughts on why?</p>
<p>Thanks,<br>
Epoching</p>

username: albanD
cooked: <p>In your first case, you actually modify <code>targets</code> inplace. Which is what the dataloader returned. It is very possible that you are incrementally changing your whole Dataset by moving each target to the GPU one by one.</p>

username: Epoching
cooked: <p>Ah interesting, are you thinking that doing this in-place keeps each minibatch on the GPU (because the DataLoader is passing original references from the source Dataset)?</p>
<p>If this is the case, why doesn’t this happen for a simple image classification dataloader (made from torchvision.datasets, for example):</p>
<pre><code class="lang-auto">images = images.cuda()
labels = labels.cuda()
</code></pre>
<p>I also forgot to mention that we have to use a custom <code>collate_fn</code>, to get the DataLoader to output batches of a list of image tensors, and a list of dictionaries.</p>
<p>Maybe the default <code>collate_fn</code> argument makes copies of the Dataset examples via <code>torch.stack</code> so it doesn’t alter the device status of the original data in the Dataset.</p>
<p>Will definitely look into this suggestion tonight! Thanks for the tip <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: albanD
cooked: <p>Hi,</p>
<p>Because <code>labels = labels.cuda()</code> this does not change the original <code>labels</code> inplace. But your code does change the original dictionary inplace ! So if this dictionary is the same as the one stored in your Dataset, then you change the one stored in your dataset.</p>

username: Epoching
cooked: <p>This is correct! Below is an example of the collate_fn I was using to batch together examples from the Dataset:</p>
<pre><code class="lang-auto">def myCollater(examples):
    input_imgs = []
    target_dicts = []
    for ex in examples:
        input_imgs.append(ex[0])
        target_dicts.append(ex[1])
        
    return torch.stack(input_imgs), target_dicts
</code></pre>
<p>I ended up appending the dictionaries to a list that came directly from the Dataset. The input image tensors didn’t stay on the GPU when training because I thought that torch.stack() basically combines multiple image tensors into one single batched tensor (which is a copy of all of the original tensors, so putting it on the GPU isn’t necessarily putting the original image tensors on the GPU).</p>
<p>Also, if we wanted to train a Mask_RCNN with variable image sizes, then the collate function should return a list of image tensors (rather than a single stacked/batched image tensor, since you can’t stack multiple tensors of variable size). This meant I had to switch my return function of <code>myCollater</code> to:</p>
<p><code>return input_imgs, target_dicts</code></p>
<p>This also means that I have to now put the images on the GPU one at a time (since <code>input_imgs</code> is a list of image tensors), so in the training loop we now have:</p>
<pre><code class="lang-auto">for j in range(len(images)):
    images[j] = images[j].to(device)
</code></pre>
<p>instead of:</p>
<pre><code class="lang-auto">images = images.to(device)
</code></pre>
<p>I thought that this would now cause the images to stay on the GPU, since they are supplied directly from the Dataset class, and appended in a list through the <code>myCollater</code> function. Turns out that these images don’t stay on the GPU, and everything runs fine.</p>
<p>I ran a small experiment to find out that when you add a tensor to a list, a copy of the tensor is made every time it’s added to a list. (So each list has it’s own unique reference to a tensor). For dictionary values, it’s the same thing!</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/6/96609dc7180f6a0f1293ff4f7d88a4552003e8b5.png" data-download-href="https://discuss.pytorch.org/uploads/default/96609dc7180f6a0f1293ff4f7d88a4552003e8b5" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/6/96609dc7180f6a0f1293ff4f7d88a4552003e8b5_2_690x392.png" alt="image" data-base62-sha1="lsiCrzBfeRmrW3rW3GZJEcZN1ml" width="690" height="392" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/6/96609dc7180f6a0f1293ff4f7d88a4552003e8b5_2_690x392.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/6/96609dc7180f6a0f1293ff4f7d88a4552003e8b5_2_1035x588.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/6/96609dc7180f6a0f1293ff4f7d88a4552003e8b5.png 2x" data-dominant-color="F5F5F5"><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">image</span><span class="informations">1323×752 125 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>So my GPU memory problem was caused from appending <em>original dictionaries</em> of tensors. Appending dictionaries (or lists) doesn’t necessarily copy the objects for you automatically, which is why we have to do it ourselves explicitly in the training loop (or in the collate_fn). Oops!</p>
<p>Thanks for the info again <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>