Using multiple optimizers without multiple calls to training_step, in order not to slow down the training

[amansouri3476]

Hi, thanks for your great project!

I have a question about using multiple optimizers. I'll share a minimal example below.

Let's say I have 3 modules: F, Z, M
I want to update the parameters for M upon every batch that comes into the training loop. I want to update Z every 5 batches, and F every 10 batches. I have assigned a 'name' key to each optimizer, and I store these intervals in correspondence with these names. In optimizer step, if the batch_idx is a multiple of 5 for instance, then optimizer_step is invoked to update the parameters of Z. Here's the code:

# We override this function to alternate the schedule for optimizer steps
    def optimizer_step(
        self,
        epoch,
        batch_idx,
        optimizer,
        optimizer_idx,
        optimizer_closure,
        on_tpu=False,
        using_native_amp=False,
        using_lbfgs=False,
    ):

        key = optimizer.param_groups[0]['name']

        if batch_idx % self.optimizer_intervals[key] == 0:

            # the closure (which includes the `training_step`) will be executed by `optimizer.step`
            optimizer.step(closure=optimizer_closure)

        else:

            # call the closure by itself to run `training_step` + `backward` without an optimizer step
            optimizer_closure()

I have followed the style of examples here.
When there are multiple optimizers, why is there a need for training step to be run with different optimizer_idx? Or put better, what should I do if I want the training_step to run 'just once', get that loss, and do optimizer_step based on batch_idx. A pseudocode would be like this:

training_step():

        # do some computation

        return {'loss':train_loss}

optimizer_step(
        self,
        epoch,
        batch_idx,
        optimizer,
        optimizer_idx,
):

        # takes train_loss and doesn't call training_step again, and based on batch_idx do the following
        # I need to be sure that backward() and zero_grad() are handled properly.
        
        # depending on optimizer_idx, optimizer would be one of my 3 optimizers: 'F', 'Z', 'M'
        key = optimizer.param_groups[0]['name']


        # self.optimizer_intervals = {'M':1, 'Z':5, 'F':10}
        if batch_idx % self.optimizer_intervals[key] == 0:
            optimizer.step()

        else:
            # do nothing, keep the train_loss, and continue to next optimizer


# once all 3 optimizers have called this method (optimizer_step), then the
# training should proceed to the next batch, a new train_loss be computed.
# Again all optimizers come to this method and do
# zero_grad and backward based on batch_idx

It seems using optimizer_closure, would invoke training_step again, for the next optimizer. So if you have 4 optimizers, it's like training_step is called 4 times, which slows down the program, and is not one would want. Ideally you'd want one call to training_step and here in optimizer_step, each optimizer should update the parameters on its turn.

Can I get that behaviour if I remove closure=optimizer_closure from optimizer.step argument?

I'm not sure if I was clear or not, so I'll put another psuedocode of what I have in mind for things to happen under the hood:

output = process_one_batch()


# during the whole loop, output is untouched, and won't be recomputed. backward is called on the
# same 'output' during the loop.

# compute the grads
zero_grad() for all optimizers
output.backward()

for optimizer in optimizer_list:
      
      if batch_idx % optimizer_turn == 0:
            optimizer.step()
      else:
             continue

Many thanks for your support.

[amansouri3476]

I went deeper in the docs and found that when you have multiple optimizers, here’s what happens under the hood:

for epoch in epochs:
    for batch in data:
        for opt in optimizers:
            loss = train_step(batch, batch_idx, optimizer_idx)
            opt.zero_grad()
            loss.backward()
            opt.step()

    for lr_scheduler in lr_schedulers:
        lr_scheduler.step()

So even another pseudocode for what I want under the hood, is the following:
(i.e. I don’t want the loss to be computed again with each optimizer)

for epoch in epochs:
    for batch in data:
        for opt in optimizers:
            opt.zero_grad()

        loss = train_step(batch, batch_idx)
        for opt in optimizers:
            if some_condition on opt_idx:
                 opt.step()
            else:
                 continue

    for lr_scheduler in lr_schedulers:
        lr_scheduler.step()