Tutorial "step optimizer at arbitrary intervals" does not call train_step every iteration

[lizhitwo]

🐛 Bug

The tutorial step optimizer at arbitrary intervals has a bug/feature(*) that the closure is evaluated (with optimizer_idx=1) only every other batch, not every batch.

(*bug/feature) Note: this behavior is intended for e.g. GAN training where a part of the model is only trained every n batches, and the tutorial has the expected behavior. But the language looks like the tutorial is also applicable to people who are trying to implement accumulate_grad_batches here (Pytorch-lightning forbids us using accumulate_grad_batches if we override optimizer_step). For this use case, using the tutorial leads to the unexpected behavior.

Please reproduce using the BoringModel

To Reproduce

Colab

Difference: overrode optimizer_step, added printing in training_step.

Expected behavior

The training step is called every batch.

training_step(batch_idx=0)
training_step(batch_idx=1)
training_step(batch_idx=2)
training_step(batch_idx=3)
training_step(batch_idx=4)
training_step(batch_idx=5)
training_step(batch_idx=6)
...

Actual behavior

training_step(batch_idx=1)
training_step(batch_idx=3)
training_step(batch_idx=5)
training_step(batch_idx=7)
training_step(batch_idx=9)
training_step(batch_idx=11)
training_step(batch_idx=13)
...

[awaelchli]

@lizhitwo Maybe it wasn't like this at the time of writing this issue, but we now have a warning at the top of this example:

If you are overriding this method, make sure that you pass the optimizer_closure parameter to optimizer.step() function as shown in the examples because training_step(), optimizer.zero_grad(), backward() are called in the closure function.

What do you think of adding this comment: #8521 ?

[lizhitwo]

You are right, when I opened this ticket the warning was not there.
This change clarifies things; my thoughts are:
(1) Would it be clearer if (which includes the `training_step`) changes into (including `training_step()`; see warning above)?
(2) If it were up to me, I would add a tutorial for implementing accumulate_grad_batches because it involves calling optimizer_closure() without calling optimizer.step(), which might be out of scope of this ticket? If you want to add this tutorial, it would look something like

    def optimizer_step(
        self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure,
        on_tpu=False, using_native_amp=False, using_lbfgs=False,
    ):
        if (batch_idx + 1) % 2 == 0:
            # step optimizer every two iterations
            optimizer.step(closure=optimizer_closure)
        else:
            # in other iterations, the closure (including `training_step()`, see warning above) needs to be run
            optimizer_closure()

(3) This is probably out of scope of this ticket, but the warning and this change does not tell readers what a closure is, and people typically don't know unless they know LBFGS. And for people who know, unless the doc changed elsewhere, they probably don't know that pytorch-lightning makes a optimizer_closure() wrapper for even SGD/Adam despite they don't need one. One suggestion is to clarify the warning as well:

If you are overriding this method, make sure that you pass the `optimizer_closure` parameter to `optimizer.step()` function as shown in the examples because this closure function is how Lightning calls the forward pass (`training_step()`), `optimizer.zero_grad()`, and `backward()`, which is done inside `optimizer.step()` no matter which optimizer is used.