Hi there, I’ve found that many people here getting frustrated about the CUDA memory usage in pytorch. I recently write a tool named pytorch_memlab to help people better understand what’s going on underneath the python code.
Here are some sample outputs:

line_profiler:

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()


mem_reporter:

========= 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
-------------------------------------------------------------------------------

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.
Check out the source code on Github:

  
      
      GitHub
  
  
    

Stonesjtu/pytorch-memlab 40

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


  
  
    
    
  
  


Or you can just install it on pip:
pip install pytorch_memlab