How to know total GPU memory required

[hkim716]

I'm trying to analyze total GPU memory required for my current model and dataset.
I used this way to check the memory required.

    for epoch in range(num_epochs):
        losses = []
        train_loss = train(epoch)
        test_loss = test(epoch)
        print(f'batch size:{batch_size} Hidden_Ch {hidden_channel} Lat {x_lat}')
        print('Epoch: {:02d}, Train MSE: {:.4f}, Test MSE: {:.4f}'.format(epoch, train_loss, test_loss))
        a = torch.cuda.memory_allocated(0)
        print(a)

The only message that I can see is about epoch 1, saying Epoch: 00, Train MSE: 0.2684, Test MSE: 0.2386
and the printed number was 12387380224, which is about 12.38GB.

After the Epoch 00, it gives me OOM error message.

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 23.65 GiB total capacity; 22.62 GiB already allocated; 7.56 MiB free; 22.62 GiB reserved in total by PyTorch)

I think the number 12.38GB is not very accurate because the model is trained on Quadro RTX 6000, which has 24GB of memory.
How could I estimate GPU memory correctly? Is there any other huge memory consumption that I don't even know?

[rusty1s]

Analyzing GPU memory is indeed tricky. You can take a look at the PyTorch profiler. PyG also provides some simple methods to estimate GPU costs via the torch_geometric.profile package.