Memory spike Inplace sum of two torch.tensor vector

data
1

I am trying to do an InPlace addition of two torch.tensor vectors.
If the individual sizes are in order X MB after addition it is reaching to in GBs.
I am using python’s “+=” operator doing this also tried with add_().

If someone can help me with the memory efficient way to do this it will be great help for me.

This is the code snippet:-
for x in range(10000):
x1+=func(z) #func(z) returns a tensor of same shape as x1

resulting x1 seeing a memory spike.

2

Doing either += or add_ seem fine to me.

What does f do? Do any of your tensors require gradients? It could mean you accumulating a larger and larger autograd graph over time.

3

Hi,
Thanks for your reply. Yes I was doing that. That was the reason for accumulating large memory.

4

If this is still a problem for you, one thing you could do in this case is to apply activation checkpointing torch.utils.checkpoint — PyTorch 2.3 documentation. Instead of saving the activations during forward, accumulating them over time. They will be computed as needed during backward in chunks, where the size of the chunks is as small as you need to suppress the memory spike. (the size of the chunk is determined by how much of the forward you wrap in a given checkpoint)

5

First of all Thank you so much. It was so helpful.
What if I extract or remove the grad post propagate function?

6

What do you mean by the grad post propagate function?

7

I mean from the propagate function of MessagePassing we get back the tensor with the grad_function(grad_fn=) attached to it. That increases the memory utilisation.

Instead of that, we just take the “data” part. It brings down the memory utilisation.
Is this a good practice of it can introduce some conflicts/anomalies in the result?

8

I see, if you don’t need to backward through that graph, you could do torch.no_grad to avoid creating the graph in the first place.