To introduce a memory toolkit for pytorch

username: Stone
cooked: <p>Hi there, I’ve found that many people here getting frustrated about the CUDA memory usage in pytorch. I recently write a tool named <code>pytorch_memlab</code> to help people better understand what’s going on underneath the python code.</p>
<p>Here are some sample outputs:</p>
<ul>
<li>line_profiler:</li>
</ul>
<pre><code class="lang-auto">Function: work at line 71

Line # Max usage   Peak usage diff max diff peak  Line Contents
===============================================================
71                                               @profile
72                                               def work():
73                                                   # comment
74   885.00K        1.00M   40.00K    0.00B          linear = torch.nn.Linear(100, 100).cuda()
75   925.00K        1.00M   40.00K    0.00B          linear_2 = torch.nn.Linear(100, 100).cuda()
76   965.00K        1.00M   40.00K    0.00B          linear_3 = torch.nn.Linear(100, 100).cuda()
</code></pre>
<ul>
<li>mem_reporter:</li>
</ul>
<pre><code class="lang-auto">========= before backward =========
Element type                                            Size  Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight                                          (1024, 1024)     4.00M
bias                                                 (1024,)     4.00K
Tensor0                                          (512, 1024)     2.00M
Tensor1                                                 (1,)   512.00B
-------------------------------------------------------------------------------
Total Tensors: 1573889  Used Memory: 6.00M
The allocated memory on cuda:0: 6.00M
-------------------------------------------------------------------------------
========= after backward =========
Element type                                            Size  Used MEM
-------------------------------------------------------------------------------
Storage on cuda:0
weight                                          (1024, 1024)     4.00M
weight.grad                                     (1024, 1024)     4.00M
bias                                                 (1024,)     4.00K
bias.grad                                            (1024,)     4.00K
Tensor0                                          (512, 1024)     2.00M
Tensor1                                                 (1,)   512.00B
-------------------------------------------------------------------------------
Total Tensors: 2623489  Used Memory: 10.01M
The allocated memory on cuda:0: 10.01M
-------------------------------------------------------------------------------
</code></pre>
<p>Since it’s a very new repo, and I’m mainly developing features I want most.  feel free to raise any issues and feature requests in this thread or on GitHub.</p>
<p>Check out the source code on Github:<br>
<aside class="onebox whitelistedgeneric">
  <header class="source">
      <img src="https://github.githubassets.com/favicon.ico" class="site-icon" width="16" height="16">
      <a href="https://github.com/Stonesjtu/pytorch-memlab" target="_blank" rel="nofollow noopener">GitHub</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars0.githubusercontent.com/u/4556044?s=400&amp;amp;v=4" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://github.com/Stonesjtu/pytorch-memlab" target="_blank" rel="nofollow noopener">Stonesjtu/pytorch-memlab</a></h3>

<p>A lab to do simple and accurate memory related experiments on pytorch - Stonesjtu/pytorch-memlab</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Or you can just install it on pip:</p>
<pre><code class="lang-auto">pip install pytorch_memlab
</code></pre>

username: yoelshoshan
cooked: <p>Thanks for sharing!<br>
I was doing some of these things manually, and it’s a great a idea to make a nice package for it. Possibly even including it in official pytorch versions in the future <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: Stone
cooked: <p>Glad you like it!<br>
I used to do these things line-by-line, manually printing the memory usage, printing the tensor size, etc. I think by building such a library I can save people’s time from adding and deleting logging codes for many scenarios.</p>
<p>There are some memories allocated by pytorch that I cannot touch using pure python. It requires some code changes to make it work natively with pytorch.</p>

username: YONG_ZHANG
cooked: <p>Thank you for sharing. I just come across with a GPU memory. Your tool is just what I want.<br>
I will use it to investigate my model.</p>