Fix an issue that caused Lightning to extract the batch size even though it was set by the user in LightningModule.log

[rohitgr7]

What does this PR do?

Fixes: #10349 (comment)
Fixes: #10576

Seems like after we added the warning to check for batch_size for an ambiguous batch, a few users realized that the batch_size used internally isn't correct. We always make a call to extract batch size even when it's set by the user explicitly inside self.log(batch_size=...) which IMO is not required because we have a recursive call to extract batch size which causes unnecessary extraction for it even when it's not required plus it will raise a warning if someone specifies it and we still look for it. So now the call will be made inside .log only when batch_size is not specified.

In short: this fixes an issue that caused Lightning to extract the batch size even though it was set by the user in LightningModule.log.

So an overview of the solution is:
track the batch and use it to extract batch_size inside self.log call when it's not specified explicitly and metric is being logged with on_epoch=True and will be averaged at the end.

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

cc @Borda @carmocca @edward-io @ananthsub

[rohitgr7]

Hi, @rohitgr7 Thanks a lot for tackling this issue! I left a few comments in code, feel free to let me know if you have any questions. A generic idea from my point of view, I think extract_data_batch() is this kind function we could convert as hook and delegate to user to implement, because they know their batch, it will be easily for them to define the batch_size or how to get batch_size. While here we had a lot assumptions about what the batch is going to look like, but we still can only extract valid batch_size under very limited situations. Anyway, I don't want to interrupt this PR discussion thread, I will initialize an issue to start the discussion. cc: @tchaton @ananthsub

if we do this, then we will have two different ways by which user can define their batch size. One will be this hook and another one will be self.log(..., batch_size=). They still can define methods within their lightning module and can call it to retrieve the batch_size and use it within self.log. Plus we already have too many hooks and this seems more like a utility method to me. Just a personal opinion :)

[tchaton]

To add to @rohitgr7 point,

self.log(..., batch_size=) was more a rescue solution and I personally don't find it elegant to provide the batch size, especially if you are logging outside of the *_step methods of the LightningModule where the current batch is unavailable. But this might affect a minority of users.

Plus we already have too many hooks
This is entirely true and we want to enforce only 1 way to do anything within Lightning. I believe we would have to weigh the hook vs the self.log(..., batch_size=) in terms of usability/simplicity.

@carmocca @awaelchli @SeanNaren Any thoughts ?

[rohitgr7]

@tchaton let's discuss more here: #10635