Is this the right way to compute gradients of two losses from two different NN's in pytorch?

Viewed 335

I have a NN defined in pytorch and I have created two instances of that net as self.actor_critic_r1 and self.actor_critic_r2. I calculate the losses of each net i.e. loss1 and loss2 and I sum it up and calculate the grads in the following way,

loss_r1 = value_loss_r1 + action_loss_r1 - dist_entropy_r1 * args.entropy_coef
loss_r2 = value_loss_r2 + action_loss_r2 - dist_entropy_r2 * args.entropy_coef
self.optimizer_r1.zero_grad()
self.optimizer_r2.zero_grad()
loss = loss_r1 + loss_r2
loss.backward()
self.optimizer_r1.step()
self.optimizer_r2.step()
clip_grad_norm_(self.actor_critic_r1.parameters(), args.max_grad_norm)
clip_grad_norm_(self.actor_critic_r2.parameters(), args.max_grad_norm)

Alternatively, should I update the loss individually like this,

self.optimizer_r1.zero_grad()
(value_loss_r1 + action_loss_r1 - dist_entropy_r1 * args.entropy_coef).backward()
self.optimizer_r1.step()
clip_grad_norm_(self.actor_critic_r1.parameters(), args.max_grad_norm)
self.optimizer_r2.zero_grad()
(value_loss_r2 + action_loss_r2 - dist_entropy_r2 * args.entropy_coef).backward()
self.optimizer_r2.step()
clip_grad_norm_(self.actor_critic_r2.parameters(), args.max_grad_norm)

I am not sure if this the right approach to update a network with multiple loss please provide your suggestion.

1 Answers

It should be the sum approach. If there is no interplay then the gradient of the 'wrong' loss will be zero for the 'wrong' optimizer anyway, and if there is interplay you likely want to optimize for that interplay.

Only if you know that there is interplay but you do not want to optimize for it should you use approach #2.

Related