auto_scale_batch_size won't reset current_epoch

[maxjeblick]

🐛 Bug

When auto_scale_batch_size is enabled, the model is initially trained with varying batch sizes. When training begins, trainer.current_epoch equals 1 instead of 0.

To Reproduce

Either observe the progress bar or use a simple callback to track the epoch number, once with auto_scale_batch_size enabled and once with auto_scale_batch_size disabled.

from pytorch_lightning import Callback

class PrintCallback(Callback):
    
    def __init__(self):
        self.observed_epochs = []
        
    def on_train_epoch_start(self, trainer, pl_module):
        print(f'Current Epoch: {trainer.current_epoch}')
        self.observed_epochs.append(trainer.current_epoch)

[rohitgr7]

since it calls it with various batch_sizes did you find where it sets current_epoch to 0 before checking it on next batch_size?

[maxjeblick]

The problem is during model checkpointing. The checkpoint sets 'epoch': self.current_epoch + 1,. That checkpoint will be loaded after having completed the batch size finder. During batch size scaling, the epoch won't be increased.

[edenlightning]

Currently blocked until trainer.tune is added.