63 lines
2.0 KiB
Python
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
|