How can we use Pytorch Autograd for sequence optimization (in a for loop)?

Question

I want to optimize a sequence in a for loop using Pytorch Autograd. I am using LBFGS .

loss = 0.0
for i in range(10):
     x = f(x,z[i])
     loss = loss + mse_loss(x,x_GT)

Say the sequence length is 10. I want to optimize x as well as z (z is a tensor array) , these are learnable parameters. Note the x will be updated in the loop.

x_GT is ground truth data.

To run this, I have to open:

loss.backward(retain_graph=True)

Is there a better way to do so (To make it run faster)?

Answer 1

The code you provided is actually perfectly fine:

loss = torch.zeros(1)
for i in range(10):
     x = f(x, z[i])
     loss += mse_loss(x, x_GT)

It will accumulate the loss over the loop steps. The backward pass only needs to be called once, though, so you are not required to retain the graph on it:

>>> loss.backward()

I don't believe not retaining the graph will make your code run any faster. It only adds to the memory load since it has to save all activations on the graph, expecting a second backward pass to come.