Why does GPU memory fluctuate when using negative sampling in PyG DataLoader (with fixed batch size and only training process)?

[wilmerwang]

Hi, I am using HAN for heterogeneous network link prediction. During training, I observed that the GPU memory usage fluctuates. I have the following questions:

1. Does using a negative sampler inside the DataLoader inherently cause GPU memory usage to fluctuate, even if batch_size is fixed?

2. Is there any way to force PyG / PyTorch to allocate the maximum GPU memory upfront, so the usage stays stable (close to max) instead of fluctuating?

Thanks.

[akihironitta]

@wilmerwang I'm not sure what your data look like, but in graph learning, it's common for input sizes to vary because each (sub)graph generally has different number of nodes and edges. I'd suggest providing a script to reproduce the behaviour you're seeing so that someone in the PyG community could possibly help you.

[speckhard]

@wilmerwang, I think you want to use dynamic batching to try to fix the memory size of the batch. Here's an article with more general information about dynamic batching. It is implemented in pytorch here, and here's a colab discussing it's use

The basic idea is what @akihironitta noted. Graphs have different numbers of nodes/edges. To efficiently batch the data, you have to pad the total number of nodes/edges in your batch. The dynamic batching method works by the user specifying a maximum number of nodes/edges, which serves as the padding target (pad until you reach that number). That way the batch has always the same number of total nodes/edges. At the moment the PyG dynamic batching won't give you constant memory usage. I am trying to make some fixes to the dynamic batching code to make it more efficient, if you have ideas/requests, please add them here:
#10466

My feeling though is that PyG dynamic batching will give you much more stable GPU usage and is worth trying.