CUDA out of memory where num_workers >= 2

shirui-japina: <p>When I train my network,<br>
it can work well when<br>
<code>num_worker = 0</code><br>
or<br>
<code>num_worker = 1</code></p>
<p>But it will<br>
<strong>CUDA out of memory</strong><br>
when<br>
<code>num_worker &gt;= 2</code> .</p>
<p>How can I solve this problem?<br>
Or to say, all I can do is to change to a better GPU only?</p>

ptrblck: <p>If you are loading the data onto the CPU (as would be the usual work flow), the number of workers should not change the usage of the GPU memory.<br>
Could you post your <code>Dataset</code> and how you are using the <code>DataLoader</code> such that your device is running out of memory?</p>

shirui-japina: <p>Thanks for your comment. <img src="https://discuss.pytorch.org/images/emoji/apple/upside_down_face.png?v=9" title=":upside_down_face:" class="emoji" alt=":upside_down_face:"><br>
I am referring <a href="https://github.com/amdegroot/ssd.pytorch" rel="nofollow noopener">https://github.com/amdegroot/ssd.pytorch</a>.</p>
<ul>
<li>
<code>Dataset</code>:</li>
</ul>
<pre><code class="lang-auto">class VOCDetection(data.Dataset):
    def __init__(self, root,
                 image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
                 transform=None, target_transform=VOCAnnotationTransform(),
                 dataset_name='VOC0712'):
        self.root = root
        self.image_set = image_sets
        self.transform = transform
        self.target_transform = target_transform
        self.name = dataset_name
        self._annopath = osp.join('%s', 'Annotations', '%s.xml')
        self._imgpath = osp.join('%s', 'JPEGImages', '%s.jpg')
        self.ids = list()
        for (year, name) in image_sets:
            rootpath = osp.join(self.root, 'VOC' + year)
            for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')):
                self.ids.append((rootpath, line.strip()))

    def __getitem__(self, index):
        im, gt, h, w = self.pull_item(index)

        return im, gt

    def __len__(self):
        return len(self.ids)

    def pull_item(self, index):
        img_id = self.ids[index]

        target = ET.parse(self._annopath % img_id).getroot()
        img = cv2.imread(self._imgpath % img_id)
        height, width, channels = img.shape

        if self.target_transform is not None:
            target = self.target_transform(target, width, height)

        if self.transform is not None:
            target = np.array(target)
            img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])
            # to rgb
            img = img[:, :, (2, 1, 0)]
            # img = img.transpose(2, 0, 1)
            target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
        return torch.from_numpy(img).permute(2, 0, 1), target, height, width
        # return torch.from_numpy(img), target, height, width

    def pull_image(self, index):
        img_id = self.ids[index]
        return cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)

    def pull_anno(self, index):
        img_id = self.ids[index]
        anno = ET.parse(self._annopath % img_id).getroot()
        gt = self.target_transform(anno, 1, 1)
        return img_id[1], gt

    def pull_tensor(self, index):
        return torch.Tensor(self.pull_image(index)).unsqueeze_(0)
</code></pre>
<ul>
<li>using the  <code>DataLoader</code>:</li>
</ul>
<pre><code class="lang-auto">    data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True)
    # create batch iterator
    batch_iterator = iter(data_loader)
    for iteration in range(args.start_iter, cfg['max_iter']):
        if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
            update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,
                            'append', epoch_size)
            # reset epoch loss counters
            loc_loss = 0
            conf_loss = 0
            epoch += 1

        if iteration in cfg['lr_steps']:
            step_index += 1
            adjust_learning_rate(optimizer, args.gamma, step_index)

        # load train data
        images, targets = next(batch_iterator)

        if args.cuda:
            images = Variable(images.cuda())
            targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
        else:
            images = Variable(images)
            targets = [Variable(ann, volatile=True) for ann in targets]
        # forward
        t0 = time.time()
        out = net(images)
        # backprop
        optimizer.zero_grad()
        loss_l, loss_c = criterion(out, targets)
        loss = loss_l + loss_c
        loss.backward()
        optimizer.step()
        t1 = time.time()
        loc_loss += loss_l.data[0]
        conf_loss += loss_c.data[0]

        if iteration % 10 == 0:
            print('timer: %.4f sec.' % (t1 - t0))
            print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')

        if args.visdom:
            update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
                            iter_plot, epoch_plot, 'append')

        if iteration != 0 and iteration % 5000 == 0:
            print('Saving state, iter:', iteration)
            torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +
                       repr(iteration) + '.pth')
    torch.save(ssd_net.state_dict(),
               args.save_folder + '' + args.dataset + '.pth')
</code></pre>

ptrblck: <p>Thanks for the code. I cannot see anything obviously wrong.</p>
<p>Could you create a simple <code>DataLoader</code> loop with <code>num_workers=0</code> and <code>num_workers&gt;=2</code> and compare the memory usage via:</p>
<pre><code class="lang-python">loader = Dataloader(dataset, num_workers=...)   # use 0 and &gt;=2 in the next run
print(torch.cuda.memory_allocated()/1024**2)

for images, targets in loader:
    images = images.cuda()
    targets = target.cuda()
    print(torch.cuda.memory_allocated()/1024**2)
</code></pre>
<p>and post the result please?</p>
<p>PS: <code>Variables</code> are deprecated since PyTorch <code>0.4</code> and the usage is <code>.data</code> is discouraged, as it might have unwanted side effects. You can use tensors instead of <code>Variables</code> and call <code>loss.detach()</code>, if you want to store the tensor without the computation graph.</p>

shirui-japina: <p>Thanks for your comment. <img src="https://discuss.pytorch.org/images/emoji/apple/upside_down_face.png?v=12" title=":upside_down_face:" class="emoji" alt=":upside_down_face:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/upside_down_face.png?v=12" title=":upside_down_face:" class="emoji" alt=":upside_down_face:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p><code>num_workers=0</code></p>
</blockquote>
<p>after<br>
<code>loader = Dataloader(dataset, num_workers=...)</code><br>
<code>torch.cuda.memory_allocated()/1024**2</code><br>
103.033203125</p>
<p>after<br>
<code>for images, targets in loader:</code><br>
<code>torch.cuda.memory_allocated()/1024**2</code><br>
119.5205078125<br>
336.71630859375<br>
336.76123046875<br>
336.79931640625<br>
…</p>
<blockquote>
<p><code>num_workers=2</code></p>
</blockquote>
<p>after<br>
<code>loader = Dataloader(dataset, num_workers=...)</code><br>
<code>torch.cuda.memory_allocated()/1024**2</code><br>
103.033203125</p>
<p>after<br>
<code>for images, targets in loader:</code><br>
<code>torch.cuda.memory_allocated()/1024**2</code><br>
119.5205078125<br>
336.71630859375<br>
336.76123046875<br>
336.79931640625<br>
…</p>
<p>And the problem I also found is, sometimes it can actually run even <code>--num_workers=2</code> or <code>--num_workers=4</code> in debug mode.</p>
<p>PyTorch version=1.5.0<br>
PS: Maybe I should have to modify the codes like:  <code>Variables</code>, <code>.data</code>.</p>

ptrblck: <p>Thanks for the debugging.<br>
It seems that the memory usage is not increased by the number of workers.</p>
<aside class="quote no-group" data-username="shirui-japina" data-post="5" data-topic="85061">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/shirui-japina/48/15285_2.png" class="avatar"> shirui-japina:</div>
<blockquote>
<p>Maybe I should have to modify the codes like: <code>Variables</code> , <code>.data</code> .</p>
</blockquote>
</aside>
<p>Yes, start with updating the code and let’s then continue to narrow down where the OOM issue is created.</p>

Mona_Jalal: <p>In my OOM case changing num_workers for the dataloader from 10 to 0 doesn’t change the situation in which OOM raises.</p>

ptrblck: <p>Good, since the GPU memory allocations should not be changed in the common use case where no <code>CUDATensors</code> are used in the <code>Dataset</code>.</p>