GPU memory that model uses

blackbirdbarber: <p>How can we simple calculate the GPU memory model (nn.Module) use? Just a single GPU unit.</p>

ptrblck: <p>To calculate the memory requirement for all parameters and buffers, you could simply sum the number of these and multiply by the element size:</p>
<pre><code class="lang-python">mem_params = sum([param.nelement()*param.element_size() for param in model.parameters()])
mem_bufs = sum([buf.nelement()*buf.element_size() for buf in model.buffers()])
mem = mem_params + mem_bufs # in bytes
</code></pre>
<p>However, this will not include the peak memory usage for the forward and backward pass (if that’s what you are looking for).</p>

blackbirdbarber: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="56822">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>peak memory usage</p>
</blockquote>
</aside>
<p>What’s the peak memory usage?</p>

ptrblck: <p>During training you are using intermediate tensors needed to backpropagate and calculate the gradients. These intermediate tensors will be freed once the gradients were calculated (and you haven’t used <code>retain_graph=True</code>), so you’ll see more memory usage during training than the initial model parameters and buffers would use.</p>

vainaijr: <p>Is peak memory usage equivalent to forward/backward pass size here?</p>
<pre><code class="lang-auto">----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------
</code></pre>

ptrblck: <p>It might be, but I’m not sure which utility you are using and how it estimates the memory usage.</p>

blackbirdbarber: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, How can we measure the peak memory usage? This  sounds like the most important question not to break into CUDA out of memory errors. Should I ask the separate question for this?</p>

ptrblck: <p><code>torch.Cuda.max_memory_allocated()</code> should give you the max value. I’m not sure, if your currently used logging library gives a matching number, but it would be interesting to see.</p>

blackbirdbarber: <p>Thanks, it was:</p>
<pre><code class="lang-auto">import torch
torch.cuda.max_memory_allocated()
</code></pre>
<p>This can help me figure out the max batch size I can use on a model, hopefully. But I wonder if something similar is present in PyTorch already.</p>
<p>However, I am not sure if this thing will also count the memory in the garbage collector that can be free after <code>gc.collect()</code>.</p>
<p>Maybe this is called <a href="https://discuss.pytorch.org/t/collecting-garbage-questions/56821">cache</a>.</p>

raining_day513: <blockquote>
<p>These intermediate tensors will be freed once the gradients were calculated (and you haven’t used <code>retain_graph=True</code> )</p>
</blockquote>
<p>Could you provide the link to the exact lines in the source? I need to investigate this part. Thank you a lot for helping people here!</p>