I have an objective function where I am looking to maximize a loss w.r.t f1 and f2, which are encoders; at the same time minimizing it w.r.t to g, which is a bijective convolution and X is just an image.
Here is how I assume it's supposed to be done.
obs_ch1, obs_ch23 = self.infomin.split_RGB_into_R_GB(obs=obs_RGB)
obs_ch1, obs_ch23 = self.infomin.R_GB_to_frame_stacked_R_GB(
obs_R=obs_ch1, obs_GB=obs_ch23)
# minimize wrt g
logits = self.infomin.compute_logits(anchor=obs_ch1, pos=obs_ch23)
loss = self.cross_entropy_loss(logits, labels)
self.infomin_discrim_optimizer.zero_grad()
loss.backward()
self.infomin_discrim_optimizer.step()
# maximize wrt f1, f2
logits = self.infomin.compute_logits(
anchor=obs_ch1.detach(), pos=obs_ch23.detach())
labels = torch.arange(logits.shape[0]).long().to(self.device)
loss = self.cross_entropy_loss(logits, labels)
self.infomin_encoders_optimizer.zero_grad()
(-loss).backward()
self.infomin_encoders_optimizer.step()
So I have both f1 and f2 using the same optimizer which is infomin_encoders_optimizer. I calculate the NCE loss w.r.t to just g then I detach that tensor and calculate NCE w.r.t f1 and f2. I take the opposite of that value and backpropagate them separately. The reason I do it this way is because I can't do them together because they are opposing objective directions. Also have to detach what comes out of g, because otherwise, it gives an error about multiple gradient updates to g.
How are you supposed to minimize and maximize an objective at the same time w.r.t to different parameters?