247 lines
7.4 KiB
Python
Executable File
247 lines
7.4 KiB
Python
Executable File
from copy import deepcopy
|
|
import pytest
|
|
import torch
|
|
|
|
from pytorch_lightning import Trainer
|
|
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
|
from tests.base import EvalModelTemplate
|
|
from tests.base.datamodules import TrialMNISTDataModule
|
|
|
|
|
|
def test_error_on_more_than_1_optimizer(tmpdir):
|
|
""" Check that error is thrown when more than 1 optimizer is passed """
|
|
|
|
model = EvalModelTemplate()
|
|
model.configure_optimizers = model.configure_optimizers__multiple_schedulers
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=1,
|
|
)
|
|
|
|
with pytest.raises(MisconfigurationException):
|
|
trainer.tuner.lr_find(model)
|
|
|
|
|
|
def test_model_reset_correctly(tmpdir):
|
|
""" Check that model weights are correctly reset after lr_find() """
|
|
|
|
model = EvalModelTemplate()
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=1,
|
|
)
|
|
|
|
before_state_dict = deepcopy(model.state_dict())
|
|
|
|
_ = trainer.tuner.lr_find(model, num_training=5)
|
|
|
|
after_state_dict = model.state_dict()
|
|
|
|
for key in before_state_dict.keys():
|
|
assert torch.all(torch.eq(before_state_dict[key], after_state_dict[key])), \
|
|
'Model was not reset correctly after learning rate finder'
|
|
|
|
|
|
def test_trainer_reset_correctly(tmpdir):
|
|
""" Check that all trainer parameters are reset correctly after lr_find() """
|
|
|
|
model = EvalModelTemplate()
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=1,
|
|
)
|
|
|
|
changed_attributes = ['callbacks', 'logger', 'max_steps', 'auto_lr_find',
|
|
'early_stop_callback', 'accumulate_grad_batches',
|
|
'checkpoint_callback']
|
|
attributes_before = {}
|
|
for ca in changed_attributes:
|
|
attributes_before[ca] = getattr(trainer, ca)
|
|
|
|
_ = trainer.tuner.lr_find(model, num_training=5)
|
|
|
|
attributes_after = {}
|
|
for ca in changed_attributes:
|
|
attributes_after[ca] = getattr(trainer, ca)
|
|
|
|
for key in changed_attributes:
|
|
assert attributes_before[key] == attributes_after[key], \
|
|
f'Attribute {key} was not reset correctly after learning rate finder'
|
|
|
|
|
|
@pytest.mark.parametrize('use_hparams', [False, True])
|
|
def test_trainer_arg_bool(tmpdir, use_hparams):
|
|
""" Test that setting trainer arg to bool works """
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
before_lr = hparams.get('learning_rate')
|
|
if use_hparams:
|
|
del model.learning_rate
|
|
model.configure_optimizers = model.configure_optimizers__lr_from_hparams
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=2,
|
|
auto_lr_find=True,
|
|
)
|
|
|
|
trainer.tune(model)
|
|
if use_hparams:
|
|
after_lr = model.hparams.learning_rate
|
|
else:
|
|
after_lr = model.learning_rate
|
|
|
|
assert before_lr != after_lr, \
|
|
'Learning rate was not altered after running learning rate finder'
|
|
|
|
|
|
@pytest.mark.parametrize('use_hparams', [False, True])
|
|
def test_trainer_arg_str(tmpdir, use_hparams):
|
|
""" Test that setting trainer arg to string works """
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
model.my_fancy_lr = 1.0 # update with non-standard field
|
|
model.hparams['my_fancy_lr'] = 1.0
|
|
before_lr = model.my_fancy_lr
|
|
if use_hparams:
|
|
del model.my_fancy_lr
|
|
model.configure_optimizers = model.configure_optimizers__lr_from_hparams
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=2,
|
|
auto_lr_find='my_fancy_lr',
|
|
)
|
|
|
|
trainer.tune(model)
|
|
if use_hparams:
|
|
after_lr = model.hparams.my_fancy_lr
|
|
else:
|
|
after_lr = model.my_fancy_lr
|
|
|
|
assert before_lr != after_lr, \
|
|
'Learning rate was not altered after running learning rate finder'
|
|
|
|
|
|
@pytest.mark.parametrize('optimizer', ['Adam', 'Adagrad'])
|
|
def test_call_to_trainer_method(tmpdir, optimizer):
|
|
""" Test that directly calling the trainer method works """
|
|
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
if optimizer == 'adagrad':
|
|
model.configure_optimizers = model.configure_optimizers__adagrad
|
|
|
|
before_lr = hparams.get('learning_rate')
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=2,
|
|
)
|
|
|
|
lrfinder = trainer.tuner.lr_find(model, mode='linear')
|
|
after_lr = lrfinder.suggestion()
|
|
model.learning_rate = after_lr
|
|
trainer.tune(model)
|
|
|
|
assert before_lr != after_lr, \
|
|
'Learning rate was not altered after running learning rate finder'
|
|
|
|
|
|
def test_datamodule_parameter(tmpdir):
|
|
""" Test that the datamodule parameter works """
|
|
|
|
# trial datamodule
|
|
dm = TrialMNISTDataModule(tmpdir)
|
|
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
|
|
before_lr = hparams.get('learning_rate')
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=2,
|
|
)
|
|
|
|
lrfinder = trainer.tuner.lr_find(model, datamodule=dm)
|
|
after_lr = lrfinder.suggestion()
|
|
model.learning_rate = after_lr
|
|
|
|
assert before_lr != after_lr, \
|
|
'Learning rate was not altered after running learning rate finder'
|
|
|
|
|
|
def test_accumulation_and_early_stopping(tmpdir):
|
|
""" Test that early stopping of learning rate finder works, and that
|
|
accumulation also works for this feature """
|
|
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
|
|
before_lr = hparams.get('learning_rate')
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
accumulate_grad_batches=2,
|
|
)
|
|
|
|
lrfinder = trainer.tuner.lr_find(model, early_stop_threshold=None)
|
|
after_lr = lrfinder.suggestion()
|
|
|
|
assert before_lr != after_lr, \
|
|
'Learning rate was not altered after running learning rate finder'
|
|
assert len(lrfinder.results['lr']) == 99, \
|
|
'Early stopping for learning rate finder did not work'
|
|
assert lrfinder._total_batch_idx == 99 * 2, \
|
|
'Accumulation parameter did not work'
|
|
|
|
|
|
def test_suggestion_parameters_work(tmpdir):
|
|
""" Test that default skipping does not alter results in basic case """
|
|
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=3,
|
|
)
|
|
|
|
lrfinder = trainer.tuner.lr_find(model)
|
|
lr1 = lrfinder.suggestion(skip_begin=10) # default
|
|
lr2 = lrfinder.suggestion(skip_begin=80) # way too high, should have an impact
|
|
|
|
assert lr1 != lr2, \
|
|
'Skipping parameter did not influence learning rate'
|
|
|
|
|
|
def test_suggestion_with_non_finite_values(tmpdir):
|
|
""" Test that non-finite values does not alter results """
|
|
|
|
hparams = EvalModelTemplate.get_default_hparams()
|
|
model = EvalModelTemplate(**hparams)
|
|
|
|
# logger file to get meta
|
|
trainer = Trainer(
|
|
default_root_dir=tmpdir,
|
|
max_epochs=3,
|
|
)
|
|
|
|
lrfinder = trainer.tuner.lr_find(model)
|
|
before_lr = lrfinder.suggestion()
|
|
lrfinder.results['loss'][-1] = float('nan')
|
|
after_lr = lrfinder.suggestion()
|
|
|
|
assert before_lr == after_lr, \
|
|
'Learning rate was altered because of non-finite loss values'
|