Let’s say I have some computational graph - is it possible to determine the amount of memory it needs? I know there is https://pytorch.org/docs/stable/cuda.html#torch.cuda.memory_allocated 5 but that only allows you monitor the total memory used (and only if you use a GPU), so that is not very convenient if you have to monitor multiple graphs at once. I imagine it somehow must be possible to recurse down the graph to get the total size?
I created a small toy example:
import torch
import torch.nn as nn
net = torch.nn.Sequential(
    nn.Conv2d(1, 1, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Sigmoid(),
    nn.MaxPool2d(kernel_size=2, stride=1, padding=0),
    nn.Conv2d(1, 1, kernel_size=3, padding=1),
    nn.LeakyReLU()
)
optim = torch.optim.SGD(params=net.parameters(), lr=0.001)
x = torch.randn((1, 1, 64, 64))
y = net(x).sum() #how much memory?
y.backward() #how much memory?