lightning/tests/callbacks/test_checkpoint_callback_fr...

63 lines
2.0 KiB
Python

import os
from pytorch_lightning import Trainer, seed_everything
from tests.base import EvalModelTemplate
def test_mc_called_on_fastdevrun(tmpdir):
seed_everything(1234)
os.environ['PL_DEV_DEBUG'] = '1'
train_val_step_model = EvalModelTemplate()
# fast dev run = called once
# train loop only, dict, eval result
trainer = Trainer(fast_dev_run=True)
trainer.fit(train_val_step_model)
# checkpoint should have been called once with fast dev run
assert len(trainer.dev_debugger.checkpoint_callback_history) == 1
# -----------------------
# also called once with no val step
# -----------------------
train_step_only_model = EvalModelTemplate()
train_step_only_model.validation_step = None
# fast dev run = called once
# train loop only, dict, eval result
trainer = Trainer(fast_dev_run=True)
trainer.fit(train_step_only_model)
# make sure only training step was called
assert train_step_only_model.training_step_called
assert not train_step_only_model.validation_step_called
assert not train_step_only_model.test_step_called
# checkpoint should have been called once with fast dev run
assert len(trainer.dev_debugger.checkpoint_callback_history) == 1
def test_mc_called(tmpdir):
seed_everything(1234)
os.environ['PL_DEV_DEBUG'] = '1'
# -----------------
# TRAIN LOOP ONLY
# -----------------
train_step_only_model = EvalModelTemplate()
train_step_only_model.validation_step = None
# no callback
trainer = Trainer(max_epochs=3, checkpoint_callback=False)
trainer.fit(train_step_only_model)
assert len(trainer.dev_debugger.checkpoint_callback_history) == 0
# -----------------
# TRAIN + VAL LOOP ONLY
# -----------------
val_train_model = EvalModelTemplate()
# no callback
trainer = Trainer(max_epochs=3, checkpoint_callback=False)
trainer.fit(val_train_model)
assert len(trainer.dev_debugger.checkpoint_callback_history) == 0