How to calculate the used GPU memory for each part as in the paper?

[liming-ai]

Hi @QipengGuo @KaiLv69 @ayyyq

Thanks for the nice work, I am wondering how to calculate the detailed used GPU memory as illustrated in the paper, such as the results in the Table 1. What tools did you use for the calculations?

[KaiLv69]

Hi, we use gc (https://gist.github.com/dojoteef/26cd46f7cc38b38e6f443c5f62411aa3) and torch.cuda.memory_reserved() to profile memory usage.

[liming-ai]

Hi, we use gc (https://gist.github.com/dojoteef/26cd46f7cc38b38e6f443c5f62411aa3) and torch.cuda.memory_reserved() to profile memory usage.

HI, @KaiLv69

I tried to use the scripts that you provided above. However, the code example you provided cannot run successfully.

My example code is:

from torch import nn
import torch
import profile, sys, threading

model = nn.Linear(20, 30).cuda()
criterion = nn.MSELoss().cuda()

memory_profiler = profile.CUDAMemoryProfiler(
    [model, criterion],
    filename='cuda_memory.profile'
)

sys.settrace(memory_profiler)
threading.settrace(memory_profiler)

inputs = torch.randn(1, 20, requires_grad=True).cuda()
output = model(inputs)
target = torch.ones(1, 30).cuda()

loss = criterion(output, target)

then I run the python file by python3 example.py, and there is the detailed log:

/home/tiger/.local/lib/python3.9/site-packages/torch/cuda/memory.py:416: FutureWarning: torch.cuda.memory_cached has been renamed to torch.cuda.memory_reserved
  warnings.warn(
/home/tiger/.local/lib/python3.9/site-packages/torch/cuda/memory.py:424: FutureWarning: torch.cuda.max_memory_cached has been renamed to torch.cuda.max_memory_reserved
  warnings.warn(
/home/tiger/.local/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py:293: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead
  warnings.warn(
Exception ignored in: <function Library.__del__ at 0x7fea20392550>
Traceback (most recent call last):
  File "/home/tiger/.local/lib/python3.9/site-packages/torch/library.py", line 131, in __del__
  File "/home/tiger/code/diffusers/examples/text_to_image/profile.py", line 134, in __call__
TypeError: 'NoneType' object is not callable

Many people may not be familiar with the use of these tools, if you can give some more detailed examples we would appreciate it very much.