Simultaneously train two model in each epoch

autograd
1 
def U_closure(self):
    self.optimizer_U_model.zero_grad()
    loss = self.cost_function()
    loss.backward(retain_graph=True)
    return loss
def E_closure(self):
    self.optimizer_E_model.zero_grad()
    loss = self.cost_function()
    loss.backward(retain_graph=True)
    return loss
# Training function
def train(self, epochs, opt_func=torch.optim.Adam):
    torch.autograd.set_detect_anomaly(True)
    self.optimizer_U_model = opt_func(self.U_model.parameters())
    self.optimizer_E_model = opt_func(self.E_model.parameters())
    for epoch in range(epochs):
        self.optimizer_U_model.step(self.U_closure)
        self.optimizer_E_model.step(self.E_closure

There are two models U_model and E_model which are needed to be trained in each epoch. But when sending closure function to optimizer it is calculating loss two times. I have same loss for both model. Should I have only one optimizer? If I have only one optimizer then can it take parameters of both modes?

2

You can use one optimiser that has access to both models’ parameters. Concerning the loss, just sum the two, the gradient won’t change.

3

What should be the parameter for line

self.optimizer = opt_func(self.U_model.parameters())

,i.e., how can I give two weights parameter in single opttimiser?

4

I was just wondering if you have tried this.

5

I guess that @mxahan has provided you with the answer. In any case, the PyTorch documentation already provides you with examples that I think can be extended to multiple models.

https://pytorch.org/docs/stable/optim.html#how-to-use-an-optimizer

In particular I am referring to this:

optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)