Understanding LightningModule hooks #15177
Replies: 1 comment 6 replies
-
|
I'm not sure if I understand your question correctly. I had a look at your notebook, but the one training with automatic optimization failed due to the following error: automatic optimization debug logs and an error |
Beta Was this translation helpful? Give feedback.
-
|
Sorry for the confusion, I meant automatic optimization in my comment above. (Updated my comment and read your questions again!)
Great catch! In short, the reason behind When a normal PyTorch code looks like the following execution: loss = training_step()
optimizer.zero_grad()
loss.backward()
optimizer.step(), you can also write it as follows using the closure function: def closure():
loss = training_step()
optimizer.zero_grad()
loss.backward()
return loss
optimizer.step(closure)With automatic optimization, Lightning uses the latter so that it's compatible with optimizers that require reevaluation of loss, such as However, with manual optimization, it's the user's responsibility to call these hooks (in whatever order they like), and thus the execution order of manual optimization cannot be the same as that of automatic optimization in your notebook. As you point out, the hook execution order in the docs is not 100% accurate, but (I think) it's documented in that way because both are basically equivalent (and also because most people are not familiar with calling
|
Beta Was this translation helpful? Give feedback.
-
|
Hello, But what about |
Beta Was this translation helpful? Give feedback.
-
|
No, it's not deprecated. It should be called at the end of the closure and before updating the weights: def closure():
loss = training_step()
optimizer.zero_grad()
loss.backward()
on_before_optimizer_step()
return loss
def optimizer_step(closure):
loss = closure()
...
update_parameters()
optimizer_step(closure)In case you're interested in digging into the codebase, here's where Lightning calls the |
Beta Was this translation helpful? Give feedback.
-
|
I see your notebook failed with |
Beta Was this translation helpful? Give feedback.
-
|
Oh, now I see! |
Beta Was this translation helpful? Give feedback.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
-
Reading from the official doc, it looks like the fit loop should be:
But I am failing to see how this is true. You can check this notebook.
There, I create two models. One with manual optimization:
And one with automatic:
Both models have one
printin every hook, so that I can check the order they are executed and they aren't the same... Is there any in-depth documentation that explains what's going on?For the manual model:
For the automatic one:
Very weirdly, notice that
on_before_optimizer_stepis never called andoptimizer_stepis called beforetraining_step, but the official doc says otherwise...Beta Was this translation helpful? Give feedback.
All reactions