How to delete all the gradients in between operations

[malfonsoarquimea]

Hi! I am currently working on a project where, for a given trained model, I perform inference batches of inputs and compute (and store) the gradients of the output with respect to the inputs.
My code is something like this

def test_step(self, batch,batch_idx):
   with torch.set_grad_enabled(True):
      gradients_list=[]
      for batch_of_inputs in batches:
         batch_of_inputs.requires_grad_()
         output=self(batch_of_inputs)
         gradients = torch.autograd.grad(
            outputs=output,
            inputs=batch_of_inputs,
            grad_outputs=torch.ones_like(output),
            retain_graph=False,
           )
         gradients_list.append(gradients.detach_())

The thing is that the used memory increases and increases until OOM error rises. I have tried to use
del gradients,batch_of_inputs,output
and the problem persisted.
What would you suggest?
Thanks in advance

[rohitgr7]

The thing is that the used memory increases and increases until OOM error rises

if the OOM error is on GPU, you can move your gradients from GPU to CPU while storing them to release some GPU memory.

[malfonsoarquimea]

First of all, you are right, it was a mistake and it is actually y output=self(batch_of_inputs) because the model is the actual lightningmodule.
On the other hand, I already tried model.zero_grad(set_to_none=True) or, in this case, self.zero_grad(set_to_none=True) but it is not working, as the gradients that I am storing are gradients of the inputs with respect to the outputs, so the parameters of the model do not require ore have any gradients (unless I am understanding it all wrong)

[rohitgr7]

you have set with torch.set_grad_enabled(True) which might compute gradients for model weights since you are computing gradients for input wrt the output.

but if it's not working, then something else is causing the memory issues.

[malfonsoarquimea]

Hi! I finally managed to make it work by adding not only model.zero_grad(set_to_none=True) but also model.requires_grad_(False) at the end of each iteration. To reduce vram usage, I also move the gradients_list to CPU

[rohitgr7]

cool, but it's still confusing why memory keeps increasing rather than being constant across testing.
do you see this behavior after model.requires_grad_(False)?

[malfonsoarquimea]

after model.requires_grad_(False), model.zero_grad(set_to_none=True) and moving the output to the CPU as it is being generated, the GPU memory usage remains constant