Fix self.log(on_epoch=True, reduce_fx=sum) on_batch_start

[carmocca]

What does this PR do?

Minimal repro:

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer

from pytorch_lightning.callbacks import Callback


class ACallback(Callback):

    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        print(trainer.callback_metrics)

    def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        print(trainer.sanity_checking, trainer.callback_metrics)

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        pl_module.log("on_train_batch_start", 1.0, reduce_fx="sum")

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        pl_module.log("on_train_batch_end", 1.0, reduce_fx="sum")

    def on_validation_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        print("on_validation_batch_start batch", batch_idx)
        #print("BEFORE", trainer._results.get("on_validation_batch_start.on_validation_batch_start"))
        pl_module.log("on_validation_batch_start", 1.0, reduce_fx="mean")
        #print("AFTER", trainer._results.get("on_validation_batch_start.on_validation_batch_start"))

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        print("on_validation_batch_end batch", batch_idx)
        #print("BEFORE", trainer._results.get("on_validation_batch_start.on_validation_batch_start"))
        pl_module.log("on_validation_batch_end", 1.0, reduce_fx="sum")
        #print("AFTER", trainer._results.get("on_validation_batch_start.on_validation_batch_start"))


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        ...

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        progress_bar_refresh_rate=0,
        limit_train_batches=3,
        limit_val_batches=3,
        max_epochs=1,
        callbacks=[ACallback()],
        log_every_n_steps=1,
    )
    trainer.fit(model, train_dataloader=train_data, val_dataloaders=val_data)


if __name__ == "__main__":
    run()

which prints

{'on_validation_batch_start': tensor(1.), 'on_validation_batch_end': tensor(6.), 'on_train_batch_start': tensor(3.), 'on_train_batch_end': tensor(3.)}
# should be
{'on_validation_batch_start': tensor(1.), 'on_validation_batch_end': tensor(3.), 'on_train_batch_start': tensor(3.), 'on_train_batch_end': tensor(3.)}

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

None

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

• 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

[github-actions]

Build Error! No Linked Issue found. Please link an issue or mention it in the body using #<issue_id>

[tchaton]

LGMT!

[rohitgr7]

does it needs a chlog?

[Borda]

lgtm