PyTorch GPU memory allocation issues (GiB reserved in total by PyTorch)

Capo_Mestre: <p>Hello,<br>
I have defined a densenet architecture in PyTorch to use it on training data consisting of 15000 samples of 128x128 images. Here is the code:</p>
<pre><code class="lang-auto">class Dense_Block(nn.Module):
    def __init__(self, in_channels):
        super(Dense_Block, self).__init__()

        self.relu = nn.ReLU(inplace = True)
        self.bn = nn.BatchNorm2d(num_features = in_channels)

        self.conv1 = nn.Conv2d(in_channels = in_channels, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv3 = nn.Conv2d(in_channels = 64, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv4 = nn.Conv2d(in_channels = 96, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)
        self.conv5 = nn.Conv2d(in_channels = 128, out_channels = 32, kernel_size = 3, stride = 1, padding = 1)

    def forward(self, x):

        bn = self.bn(x)
        conv1 = self.relu(self.conv1(bn))

        conv2 = self.relu(self.conv2(conv1))
        c2_dense = self.relu(torch.cat([conv1, conv2], 1))

        conv3 = self.relu(self.conv3(c2_dense))
        c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

        conv4 = self.relu(self.conv4(c3_dense))
        c4_dense = self.relu(torch.cat([conv1, conv2, conv3, conv4], 1))

        conv5 = self.relu(self.conv5(c4_dense))
        c5_dense = self.relu(torch.cat([conv1, conv2, conv3, conv4, conv5], 1))

        return c5_dense

class Transition_Layer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Transition_Layer, self).__init__()

        self.relu = nn.ReLU(inplace = True)
        self.bn = nn.BatchNorm2d(num_features = out_channels)
        self.conv = nn.Conv2d(in_channels = in_channels, out_channels = out_channels, kernel_size = 1, bias = False)
        self.avg_pool = nn.AvgPool2d(kernel_size = 2, stride = 2, padding = 0)

    def forward(self, x):

        bn = self.bn(self.relu(self.conv(x)))
        out = self.avg_pool(bn)

        return out

class DenseNet(nn.Module):
    def __init__(self, nr_classes):
        super(DenseNet, self).__init__()

        self.lowconv = nn.Conv2d(in_channels = 1, out_channels = 64, kernel_size = 7, padding = 3, bias = False)
        self.relu = nn.ReLU()

        # Make Dense Blocks
        self.denseblock1 = self._make_dense_block(Dense_Block, 64)
        self.denseblock2 = self._make_dense_block(Dense_Block, 128)
        self.denseblock3 = self._make_dense_block(Dense_Block, 128)

        # Make transition Layers
        self.transitionLayer1 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 128)
        self.transitionLayer2 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 128)
        self.transitionLayer3 = self._make_transition_layer(Transition_Layer, in_channels = 160, out_channels = 64)

        # Classifier
        self.bn = nn.BatchNorm2d(num_features = 64)
        self.pre_classifier = nn.Linear(64*16*16, 512)
        self.classifier = nn.Linear(512, nr_classes)

    def _make_dense_block(self, block, in_channels):
        layers = []
        layers.append(block(in_channels))
        return nn.Sequential(*layers)

    def _make_transition_layer(self, layer, in_channels, out_channels):
        modules = []
        modules.append(layer(in_channels, out_channels))
        return nn.Sequential(*modules)

    def forward(self, x):
        out = self.relu(self.lowconv(x))

        out = self.denseblock1(out)
        out = self.transitionLayer1(out)

        out = self.denseblock2(out)
        out = self.transitionLayer2(out)

        out = self.denseblock3(out)
        out = self.transitionLayer3(out)

        out = self.bn(out)
#         print(out.shape)
        out = out.reshape(-1, 64*16*16)

        out = self.pre_classifier(out)
        out = self.classifier(out)

        return out
</code></pre>
<p>Then I define my DataSet class:</p>
<pre><code class="lang-auto">class MyDataset(Dataset):
    def __init__(self, images, n, labels=None, transforms=None):
        self.X = images
        self.y = labels
        self.n = n
        self.transforms = transforms
         
    def __len__(self):
        return (len(self.X))
    
    def __getitem__(self, i):
        data = self.X.iloc[i, :]
        data = np.asarray(data).astype(np.float).reshape(1,n,n)
        
        if self.transforms:
            data = self.transforms(data).reshape(1,n,n)
            
        if self.y is not None:
            y = self.y.iloc[i,:]
            y = np.asarray(y).astype(np.float).reshape(128,) # for 128-vector of labels
            return (data, y)
        else:
            return data
</code></pre>
<p>Then I create the instances of the train, dev, and test data:</p>
<pre><code class="lang-auto">train_data = MyDataset(train_images, n, train_labels, None)
dev_data = MyDataset(dev_images, n, dev_labels, None)
test_data = MyDataset(test_images, n, test_labels, None)
</code></pre>
<p>The shapes of  <code>train_images</code> ,  <code>dev_images</code>  and  <code>test_images</code>  are respectively  <code>(15000, 16384)</code> ,  <code>(4000, 16384)</code>  and  <code>(1000, 16384)</code> . So there are in total 20000 samples of 128x128 (=16384) images.</p>
<p>The shapes of  <code>train_labels</code> ,  <code>dev_labels</code>  and  <code>test_labels</code>  are respectively  <code>(15000, 128)</code> ,  <code>(4000, 128)</code>  and  <code>(1000, 128)</code> . So there are in total 20000 samples of 128 vectors.</p>
<p>I define also a custom loss function:</p>
<pre><code class="lang-auto"># Huber
def Huber(yHat,y,delta=1.):
    n_samples = yHat.size()[0]
    n_points = yHat.size()[1]
    preds = yHat
    labels = y
    size = yHat.size()[0]*yHat.size()[1]
    diff = yHat - y
    return torch.sum(torch.where(torch.abs(diff) &lt; delta,.5*diff**2 , delta*(torch.abs(diff)-.5*delta**2))) / size
</code></pre>
<p>Then I create an instance of the model:</p>
<pre><code class="lang-auto">densenet = DenseNet(nr_classes=128).float().to('cuda:0')
</code></pre>
<p>Then I initialize parameters, create train- and dev-set dataloaders, and train the model using Adam optimizer and Huber loss-function:</p>
<pre><code class="lang-auto">def main():
    nn.init.kaiming_uniform_(list(densenet.parameters())[0], nonlinearity = 'relu')
    loader = DataLoader(train_data,batch_size=200,shuffle=False,num_workers=0)
    loader_dev = DataLoader(dev_data,batch_size=10,shuffle=None,num_workers=0)
    N_epochs = 10
    for epoch in range(N_epochs):
          optimizer = optim.Adam(densenet.parameters(), lr=.001, betas=(0.9, 0.999), eps=1e-08)
          for batch in loader:
                images = batch[0].float().to('cuda:0')
                labels = batch[1].float().to('cuda:0')
                preds = densenet(images)
                loss = Loss(preds,labels).Huber()
            
                with torch.no_grad():
                    loss_dev = 0
                    for batch_dev in loader_dev:
                        images_dev = batch_dev[0].float().to('cuda:0')
                        labels_dev = batch_dev[1].float().to('cuda:0')
                        preds_dev = densenet(images_dev)
                        loss_ = Loss(preds_dev,labels_dev).Huber()
                        loss_dev += loss_

                 optimizer.zero_grad()
                 loss.backward()
                 optimizer.step()

if __name__ == '__main__':
    multiprocessing.freeze_support()
    main()
</code></pre>
<p>I have two GPUs identified as <code>cuda:0</code> (~ 24 Gb memory) and <code>cuda:1</code> (~ 6 Gb memory)<br>
With <code>cuda:1</code> device I get the error message:</p>
<pre><code class="lang-auto">
  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1127, in &lt;module&gt;
    main()

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1074, in main
    preds = network(images) # Pass Batch

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 752, in forward
    out = self.denseblock1(out)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\container.py", line 117, in forward
    input = module(input)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 691, in forward
    c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

RuntimeError: CUDA out of memory. Tried to allocate 1.17 GiB (GPU 1; 6.00 GiB total capacity; 4.34 GiB already allocated; 16.62 MiB free; 4.34 GiB reserved in total by PyTorch)
</code></pre>
<p>Then I tried to switch to <code>cuda:0</code> device which has much more memory capacity, and the error in this case reads as:</p>
<pre><code class="lang-auto">
  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1127, in &lt;module&gt;
    main()

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 1074, in main
    preds = network(images) # Pass Batch

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 752, in forward
    out = self.denseblock1(out)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\container.py", line 117, in forward
    input = module(input)

  File "C:\Users\Admin\.conda\envs\pytorch_env\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)

  File "D:\Jupiter_playground\fashion_mnist_tidied.py", line 691, in forward
    c3_dense = self.relu(torch.cat([conv1, conv2, conv3], 1))

RuntimeError: CUDA out of memory. Tried to allocate 1.17 GiB (GPU 0; 24.00 GiB total capacity; 21.59 GiB already allocated; 372.94 MiB free; 21.69 GiB reserved in total by PyTorch)
</code></pre>
<p>Why does PyTorch allocate almost all available memory?</p>
<p>However, when I use train-set of 6 images and dev-set of 3 images (test-set of 1 image), training with <code>cuda</code>-devices works fine.</p>

JuanFMontesinos: <p>Because ur batch-size is 200 and densenet is demanding.  Why don’t you use smaller BS and increase up the the gpu memory?<br>
I mean, according to the code your batch size is 200.</p>
<p>Densenet is heavy (densily conected)</p>

Capo_Mestre: <aside class="quote no-group" data-username="JuanFMontesinos" data-post="2" data-topic="93134">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/48/76115_2.png" class="avatar"> JuanFMontesinos:</div>
<blockquote>
<p>increase up the the gpu memory</p>
</blockquote>
</aside>
<p>Thanks for your reply. What do you mean by increasing the gpu memory? Buying a new one?</p>

JuanFMontesinos: <p>Increasing the memory demand. I mean, start from a small batch size and increase it checking gpu usage with nvidia-smi. You need to compute the backward as well since it requires memory. When you see you can compute forward and backward reaching the memory limit thats the maximum batch siD you can use</p>

Capo_Mestre: <p>Thanks for the tip! I guess this is a good systematic approach <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>