Restore trainer.current_epoch after tuning (#7434)

* Add a test * Save and restore current_epoch * Update CHANGELOG * alphabetical order
Lightning-AI · May 8, 2021 · 710b144 · 710b144
1 parent 45143fd
commit 710b144
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -37,6 +37,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed DeepSpeed with IterableDatasets ([#7362](https://github.com/PyTorchLightning/pytorch-lightning/pull/7362))
 
 
+- Fixed `Trainer.current_epoch` not getting restored after tuning ([#7434](https://github.com/PyTorchLightning/pytorch-lightning/pull/7434))
+
+
 ## [1.3.0] - 2021-05-06
 
 ### Added

diff --git a/pytorch_lightning/tuner/lr_finder.py b/pytorch_lightning/tuner/lr_finder.py
@@ -288,6 +288,7 @@ def __lr_finder_dump_params(trainer, model):
  'logger': trainer.logger,
  'max_steps': trainer.max_steps,
  'checkpoint_callback': trainer.checkpoint_callback,
+ 'current_epoch': trainer.current_epoch,
  'configure_optimizers': model.configure_optimizers,
  }
 
@@ -297,6 +298,7 @@ def __lr_finder_restore_params(trainer, model):
  trainer.logger = trainer.__dumped_params['logger']
  trainer.callbacks = trainer.__dumped_params['callbacks']
  trainer.max_steps = trainer.__dumped_params['max_steps']
+ trainer.current_epoch = trainer.__dumped_params['current_epoch']
  model.configure_optimizers = trainer.__dumped_params['configure_optimizers']
  del trainer.__dumped_params
 

diff --git a/tests/tuner/test_lr_finder.py b/tests/tuner/test_lr_finder.py
@@ -77,7 +77,13 @@ def test_trainer_reset_correctly(tmpdir):
  )
 
  changed_attributes = [
- 'callbacks', 'logger', 'max_steps', 'auto_lr_find', 'accumulate_grad_batches', 'checkpoint_callback'
+ 'accumulate_grad_batches',
+ 'auto_lr_find',
+ 'callbacks',
+ 'checkpoint_callback',
+ 'current_epoch',
+ 'logger',
+ 'max_steps',
  ]
  expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
  trainer.tuner.lr_find(model, num_training=5)

diff --git a/tests/tuner/test_scale_batch_size.py b/tests/tuner/test_scale_batch_size.py
@@ -111,13 +111,13 @@ def test_trainer_reset_correctly(tmpdir):
  )
 
  changed_attributes = [
- 'max_steps',
- 'weights_summary',
- 'logger',
  'callbacks',
  'checkpoint_callback',
- 'limit_train_batches',
  'current_epoch',
+ 'limit_train_batches',
+ 'logger',
+ 'max_steps',
+ 'weights_summary',
  ]
  expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
  trainer.tuner.scale_batch_size(model, max_trials=5)