[Question] How to optimize two loss alternately with gradient accumulation?
I want to update a model by optimizing two loss alternately with gradient accumulation like this# Suppose gradient_accumulation is set to 2.
optimizer = optim(unet.parameters())
with accelerator.accumulate(unet):
    outputs = unet(input)
    loss1 = loss_func1(outputs)
    loss1.backward()
    optimizer.step() 
    optimizer.zero_grad()

with accelerator.accumulate(unet):
    outputs = unet(input)
    loss2 = loss_func2(outputs)
    loss2.backward()
    optimizer.step()
    optimizer.zero_grad()Is this correct? It appears from thedocumentationthataccelerator.accumulatewill normalize the loss and then backpropagate without updating the gradient until reachinggradient_accumulation_steps. My main concern is that the gradients accumulated by two different losses for the same model will affect each other.Hope to find some help here, thanks in advance.
2023-08-21T15:38:03Z
According to here, that seems like it would work fine:Multiple gradient updates with two separate losses and two classifiers sharing the same encoder - nlp - PyTorch Forums. During accumulation it’d be attached separately/referenced separately I believe, with wouldn’t cause them to entangle.The key here isdon’tcallzero_grad()before if possible. So I think we can essentially rewrite this as:# Suppose gradient_accumulation is set to 2.
optimizer = optim(unet.parameters())
with accelerator.accumulate(unet):
    outputs = unet(input)
    loss1 = loss_func1(outputs)
    loss1.backward(retain_graph=True)
    loss2 = loss_func2(outputs)
    loss2.backward()
    optimizer.step()
    optimizer.zero_grad()And this should work/be okay, the key is ensuringretain_graph=True
2023-08-22T08:54:27Z
Hi, thank you for the reply!I’m sorry I didn’t make my question clear. Optimizing two loss alternately (sequentially) meansloss2is computed w.r.t the model updated by optimizingloss1. I want to useaccelerator.accumulatedoing gradient accumulation with the pseudo code inMultiple gradient updates with two separate losses and two classifiers sharing the same encoder - #2 by srishti-git1110 - nlp - PyTorch Forums.import itertools
import torch

params = [encoder.parameters(), fc1.parameters(), fc2.parameters()]
optimizer = torch.optim.Adam(itertools.chain(*params), lr=0.01)

# optimizer.step() twice is what I exactly want.
for batch_idx, batch in dataloader_instance:
     # calculate lcce and lwd
     lcce.backward()
     optimizer.step()
     optimizer.zero_grad()

     lwd = -1 * lwd
     lwd.backward()
     
     for param in encoder.parameters():
          param.grad = -1*beta*param.grad
     optimizer.step()
     optimizer.zero_grad()However, your current code updates the model using the sum of the accumulated gradients of the sum ofloss1andloss2.
2023-08-22T08:59:25Z
By the way, is it necessary to setretain_graph=Truekeeping the compute graph w.r.t loss1? Sincetorch.Tensor.backwardaccumulates gradients in the leaves by default.
2023-09-11T16:46:11Z
hkunzhe:Sincetorch.Tensor.backwardaccumulates gradients in the leaves by default.My advice there relies on the solution posted just below that. If it truly does I don’t know, as I’m learning this like you in this fairly niche question. Perhaps also worth asking on that forum discussion as well