To introduce a memory toolkit for pytorch

autograd
1

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:

Or you can just install it on pip:

pip install pytorch_memlab
3 Likes
2

Thanks for sharing!
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 :wink:

3

Glad you like it!
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.

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.

1 Like
4

Thank you for sharing. I just come across with a GPU memory. Your tool is just what I want.
I will use it to investigate my model.