def test_lr_scheduler_with_no_actual_scheduler_raises(tmpdir):
"""Test exception when lr_scheduler dict has no scheduler."""
model = BoringModel()
model.configure_optimizers = lambda: {"optimizer": optim.Adam(model.parameters()), "lr_scheduler": {}}
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.raises(MisconfigurationException, match='The lr scheduler dict must have the key "scheduler"'):
trainer.fit(model)
def test_invalid_setup_method():
"""Test error message when `setup` method of `LightningModule` or `LightningDataModule` is not defined
correctly."""
class CustomModel(BoringModel):
def setup(self):
pass
class CustomDataModule(BoringDataModule):
def setup(self):
pass
class CustomBoringCallback(Callback):
def setup(self, pl_module, trainer):
pass
fit_kwargs = [
{"model": CustomModel(), "datamodule": BoringDataModule()},
{"model": BoringModel(), "datamodule": CustomDataModule()},
]
for kwargs in fit_kwargs:
trainer = Trainer(fast_dev_run=True)
with pytest.raises(MisconfigurationException, match="does not have a `stage` argument"):
trainer.fit(**kwargs)
trainer = Trainer(fast_dev_run=True, callbacks=[CustomBoringCallback()])
model = BoringModel()
with pytest.raises(MisconfigurationException, match="does not have a `stage` argument"):
trainer.fit(model)
def test_deepspeed_summary(tmpdir):
"""Test to ensure that the summary contains the correct values when stage 3 is enabled and that the trainer
enables the `DeepSpeedSummary` when DeepSpeed is used."""
model = BoringModel()
total_parameters = sum(x.numel() for x in model.parameters())
class TestCallback(Callback):
def on_fit_start(self, trainer: "pl.Trainer",
pl_module: "pl.LightningModule") -> None:
model_summary = DeepSpeedSummary(pl_module, max_depth=1)
assert model_summary.total_parameters == total_parameters
assert model_summary.trainable_parameters == total_parameters
# check the additional params per device
summary_data = model_summary._get_summary_data()
params_per_device = summary_data[-1][-1]
assert int(
params_per_device[0]) == (model_summary.total_parameters // 2)
trainer = Trainer(
strategy=DeepSpeedStrategy(stage=3),
default_root_dir=tmpdir,
accelerator="gpu",
fast_dev_run=True,
devices=2,
precision=16,
enable_model_summary=True,
callbacks=[TestCallback()],
)
trainer.fit(model)
def test_none_optimizer(tmpdir):
model = BoringModel()
model.configure_optimizers = lambda: None
trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, limit_val_batches=0.1, limit_train_batches=0.2)
with pytest.warns(UserWarning, match="will run with no optimizer"):
trainer.fit(model)
assert trainer.state.finished, f"Training failed with {trainer.state}"
def test_unknown_configure_optimizers_raises(tmpdir):
"""Test exception with an unsupported configure_optimizers return."""
model = BoringModel()
model.configure_optimizers = lambda: 1
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.raises(MisconfigurationException, match="Unknown configuration for model optimizers"):
trainer.fit(model)
def test_onecyclelr_with_epoch_interval_warns():
"""Test warning when a OneCycleLR is used and interval is epoch."""
model = BoringModel()
optimizer = optim.Adam(model.parameters())
lr_scheduler = {"scheduler": optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, total_steps=3)}
with pytest.warns(RuntimeWarning, match="Are you sure you didn't mean 'interval': 'step'?"):
_configure_schedulers_automatic_opt([lr_scheduler], None)
def test_poptorch_models_at_different_stages(tmpdir):
plugin = IPUStrategy()
trainer = Trainer(default_root_dir=tmpdir,
strategy=plugin,
accelerator="ipu",
devices=8)
model = BoringModel()
model.trainer = trainer
plugin.model = model
trainer.optimizers = model.configure_optimizers()[0]
trainer.state.fn = TrainerFn.FITTING
trainer.strategy.setup(trainer)
assert list(trainer.strategy.poptorch_models) == [
RunningStage.TRAINING, RunningStage.VALIDATING
]
for fn, stage in (
(TrainerFn.VALIDATING, RunningStage.VALIDATING),
(TrainerFn.TESTING, RunningStage.TESTING),
(TrainerFn.PREDICTING, RunningStage.PREDICTING),
):
trainer.state.fn = fn
trainer.state.stage = stage
trainer.strategy.setup(trainer)
assert list(trainer.strategy.poptorch_models) == [stage]
def test_hpc_restore_attempt(tmpdir):
"""Test that restore() attempts to restore the hpc_ckpt with highest priority."""
model = BoringModel()
trainer = Trainer(default_root_dir=tmpdir, max_steps=1, enable_checkpointing=False, logger=False)
trainer.fit(model)
hpc_ckpt_path = tmpdir / "hpc_ckpt_3.ckpt"
trainer.save_checkpoint(hpc_ckpt_path)
assert os.listdir(tmpdir) == ["hpc_ckpt_3.ckpt"]
# set weights to zero
for param in model.parameters():
torch.nn.init.constant_(param, 0)
# case 1: restore hpc first, no explicit resume path provided
trainer = Trainer(default_root_dir=tmpdir, max_steps=2, enable_checkpointing=False, logger=False)
trainer.fit(model)
for param in model.parameters():
assert param.abs().sum() > 0
torch.nn.init.constant_(param, 0)
# case 2: explicit resume path provided, restore hpc anyway
trainer = Trainer(default_root_dir=tmpdir, max_steps=3)
trainer.fit(model, ckpt_path="not existing")
for param in model.parameters():
assert param.abs().sum() > 0
def test_v1_8_0_lightning_module_use_amp():
model = BoringModel()
with pytest.deprecated_call(
match="`LightningModule.use_amp` was deprecated in v1.6"):
_ = model.use_amp
with pytest.deprecated_call(
match="`LightningModule.use_amp` was deprecated in v1.6"):
model.use_amp = False
def test_property_current_epoch():
"""Test that the current_epoch in LightningModule is accessible via the Trainer."""
model = BoringModel()
assert model.current_epoch == 0
trainer = Mock(current_epoch=123)
model.trainer = trainer
assert model.current_epoch == 123
def test_tensorboard_log_graph(tmpdir, example_input_array):
"""test that log graph works with both model.example_input_array and if array is passed externally."""
model = BoringModel()
if example_input_array is not None:
model.example_input_array = None
logger = TensorBoardLogger(tmpdir, log_graph=True)
logger.log_graph(model, example_input_array)
def test_wrong_configure_optimizers(tmpdir):
"""Test that an error is thrown when no `configure_optimizers()` is defined."""
trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
with pytest.raises(MisconfigurationException, match=r"No `configure_optimizers\(\)` method defined."):
model = BoringModel()
model.configure_optimizers = None
trainer.fit(model)
def test_property_global_step():
"""Test that the global_step in LightningModule is accessible via the Trainer."""
model = BoringModel()
assert model.global_step == 0
trainer = Mock(global_step=123)
model.trainer = trainer
assert model.global_step == 123
def test_property_local_rank():
"""Test that the local rank in LightningModule is accessible via the Trainer."""
model = BoringModel()
assert model.local_rank == 0
trainer = Mock(local_rank=123)
model.trainer = trainer
assert model.local_rank == 123
def test_is_overridden():
model = BoringModel()
datamodule = BoringDataModule()
# edge cases
assert not is_overridden("whatever", None)
with pytest.raises(ValueError, match="Expected a parent"):
is_overridden("whatever", object())
assert not is_overridden("whatever", model)
assert not is_overridden("whatever", model, parent=LightningDataModule)
class TestModel(BoringModel):
def foo(self):
pass
def bar(self):
return 1
with pytest.raises(ValueError,
match="The parent should define the method"):
is_overridden("foo", TestModel())
# normal usage
assert is_overridden("training_step", model)
assert is_overridden("train_dataloader", datamodule)
class WrappedModel(TestModel):
def __new__(cls, *args, **kwargs):
obj = super().__new__(cls)
obj.foo = cls.wrap(obj.foo)
obj.bar = cls.wrap(obj.bar)
return obj
@staticmethod
def wrap(fn):
@wraps(fn)
def wrapper():
fn()
return wrapper
def bar(self):
return 2
# `functools.wraps()` support
assert not is_overridden("foo", WrappedModel(), parent=TestModel)
assert is_overridden("bar", WrappedModel(), parent=TestModel)
# `Mock` support
mock = Mock(spec=BoringModel, wraps=model)
assert is_overridden("training_step", mock)
mock = Mock(spec=BoringDataModule, wraps=datamodule)
assert is_overridden("train_dataloader", mock)
# `partial` support
model.training_step = partial(model.training_step)
assert is_overridden("training_step", model)
def test_property_loggers(tmpdir):
"""Test that loggers in LightningModule is accessible via the Trainer."""
model = BoringModel()
assert model.loggers == []
logger = TensorBoardLogger(tmpdir)
trainer = Trainer(logger=logger)
model.trainer = trainer
assert model.loggers == [logger]
def test_property_logger(tmpdir):
"""Test that the logger in LightningModule is accessible via the Trainer."""
model = BoringModel()
assert model.logger is None
logger = TensorBoardLogger(tmpdir)
trainer = Mock(loggers=[logger])
model.trainer = trainer
assert model.logger == logger
def test_dm_apply_batch_transfer_handler(get_module_mock, accelerator, device):
expected_device = torch.device(device)
class CustomBatch:
def __init__(self, data):
self.samples = data[0]
self.targets = data[1]
class CurrentTestDM(LightningDataModule):
rank = 0
transfer_batch_to_device_hook_rank = None
on_before_batch_transfer_hook_rank = None
on_after_batch_transfer_hook_rank = None
def on_before_batch_transfer(self, batch, dataloader_idx):
assert dataloader_idx == 0
self.on_before_batch_transfer_hook_rank = self.rank
self.rank += 1
batch.samples += 1
return batch
def on_after_batch_transfer(self, batch, dataloader_idx):
assert dataloader_idx == 0
assert batch.samples.device == batch.targets.device == expected_device
self.on_after_batch_transfer_hook_rank = self.rank
self.rank += 1
batch.targets *= 2
return batch
def transfer_batch_to_device(self, batch, device, dataloader_idx):
assert dataloader_idx == 0
self.transfer_batch_to_device_hook_rank = self.rank
self.rank += 1
batch.samples = batch.samples.to(device)
batch.targets = batch.targets.to(device)
return batch
dm = CurrentTestDM()
model = BoringModel()
batch = CustomBatch((torch.zeros(5, 32), torch.ones(5, 1, dtype=torch.long)))
trainer = Trainer(accelerator=accelerator, devices=1)
model.trainer = trainer
# running .fit() would require us to implement custom data loaders, we mock the model reference instead
get_module_mock.return_value = model
trainer._data_connector.attach_datamodule(model, datamodule=dm)
batch_gpu = trainer.strategy.batch_to_device(batch, expected_device)
assert dm.on_before_batch_transfer_hook_rank == 0
assert dm.transfer_batch_to_device_hook_rank == 1
assert dm.on_after_batch_transfer_hook_rank == 2
assert batch_gpu.samples.device == batch_gpu.targets.device == expected_device
assert torch.allclose(batch_gpu.samples.cpu(), torch.ones(5, 32))
assert torch.allclose(batch_gpu.targets.cpu(), torch.ones(5, 1, dtype=torch.long) * 2)
def test_torchscript_device(device_str):
"""Test that scripted module is on the correct device."""
device = torch.device(device_str)
model = BoringModel().to(device)
model.example_input_array = torch.randn(5, 32)
script = model.to_torchscript()
assert next(script.parameters()).device == device
script_output = script(model.example_input_array.to(device))
assert script_output.device == device
def test_reducelronplateau_with_no_monitor_raises(tmpdir):
"""Test exception when a ReduceLROnPlateau is used with no monitor."""
model = BoringModel()
optimizer = optim.Adam(model.parameters())
model.configure_optimizers = lambda: ([optimizer], [optim.lr_scheduler.ReduceLROnPlateau(optimizer)])
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.raises(
MisconfigurationException, match="`configure_optimizers` must include a monitor when a `ReduceLROnPlateau`"
):
trainer.fit(model)
def test_torchscript_with_no_input(tmpdir):
"""Test that an error is thrown when there is no input tensor."""
model = BoringModel()
model.example_input_array = None
with pytest.raises(
ValueError,
match=
"requires either `example_inputs` or `model.example_input_array`"):
model.to_torchscript(method="trace")
def test_model_saves_on_gpu(tmpdir, accelerator):
"""Test that model saves on gpu."""
model = BoringModel()
trainer = Trainer(accelerator=accelerator, devices=1, fast_dev_run=True)
trainer.fit(model)
file_path = os.path.join(tmpdir, "model.onnx")
input_sample = torch.randn((1, 32))
model.to_onnx(file_path, input_sample)
assert os.path.isfile(file_path)
assert os.path.getsize(file_path) > 4e2
def test_model_saves_with_input_sample(tmpdir):
"""Test that ONNX model saves with input sample and size is greater than 3 MB."""
model = BoringModel()
trainer = Trainer(fast_dev_run=True)
trainer.fit(model)
file_path = os.path.join(tmpdir, "model.onnx")
input_sample = torch.randn((1, 32))
model.to_onnx(file_path, input_sample)
assert os.path.isfile(file_path)
assert os.path.getsize(file_path) > 4e2
def test_error_if_no_input(tmpdir):
"""Test that an error is thrown when there is no input tensor."""
model = BoringModel()
model.example_input_array = None
file_path = os.path.join(tmpdir, "model.onnx")
with pytest.raises(
ValueError,
match=r"Could not export to ONNX since neither `input_sample` nor"
r" `model.example_input_array` attribute is set.",
):
model.to_onnx(file_path)
def test_lr_scheduler_with_extra_keys_warns(tmpdir):
"""Test warning when lr_scheduler dict has extra keys."""
model = BoringModel()
optimizer = optim.Adam(model.parameters())
model.configure_optimizers = lambda: {
"optimizer": optimizer,
"lr_scheduler": {"scheduler": optim.lr_scheduler.StepLR(optimizer, 1), "foo": 1, "bar": 2},
}
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.warns(RuntimeWarning, match=r"Found unsupported keys in the lr scheduler dict: \{.+\}"):
trainer.fit(model)
def test_v2_0_resume_from_checkpoint_trainer_constructor(tmpdir):
# test resume_from_checkpoint still works until v2.0 deprecation
model = BoringModel()
callback = OldStatefulCallback(state=111)
trainer = Trainer(default_root_dir=tmpdir,
max_steps=1,
callbacks=[callback])
trainer.fit(model)
ckpt_path = trainer.checkpoint_callback.best_model_path
callback = OldStatefulCallback(state=222)
with pytest.deprecated_call(
match=
r"Setting `Trainer\(resume_from_checkpoint=\)` is deprecated in v1.5"
):
trainer = Trainer(default_root_dir=tmpdir,
max_steps=2,
callbacks=[callback],
resume_from_checkpoint=ckpt_path)
with pytest.deprecated_call(
match=r"trainer.resume_from_checkpoint` is deprecated in v1.5"):
_ = trainer.resume_from_checkpoint
assert trainer._checkpoint_connector.resume_checkpoint_path is None
assert trainer._checkpoint_connector.resume_from_checkpoint_fit_path == ckpt_path
trainer.validate(model=model, ckpt_path=ckpt_path)
assert callback.state == 222
assert trainer._checkpoint_connector.resume_checkpoint_path is None
assert trainer._checkpoint_connector.resume_from_checkpoint_fit_path == ckpt_path
with pytest.deprecated_call(
match=r"trainer.resume_from_checkpoint` is deprecated in v1.5"):
trainer.fit(model)
ckpt_path = trainer.checkpoint_callback.best_model_path # last `fit` replaced the `best_model_path`
assert callback.state == 111
assert trainer._checkpoint_connector.resume_checkpoint_path is None
assert trainer._checkpoint_connector.resume_from_checkpoint_fit_path is None
trainer.predict(model=model, ckpt_path=ckpt_path)
assert trainer._checkpoint_connector.resume_checkpoint_path is None
assert trainer._checkpoint_connector.resume_from_checkpoint_fit_path is None
trainer.fit(model)
assert trainer._checkpoint_connector.resume_checkpoint_path is None
assert trainer._checkpoint_connector.resume_from_checkpoint_fit_path is None
# test fit(ckpt_path=) precedence over Trainer(resume_from_checkpoint=) path
model = BoringModel()
with pytest.deprecated_call(
match=
r"Setting `Trainer\(resume_from_checkpoint=\)` is deprecated in v1.5"
):
trainer = Trainer(resume_from_checkpoint="trainer_arg_path")
with pytest.raises(
FileNotFoundError,
match=
"Checkpoint at fit_arg_ckpt_path not found. Aborting training."):
trainer.fit(model, ckpt_path="fit_arg_ckpt_path")
def test_lr_scheduler_with_unknown_interval_raises(tmpdir):
"""Test exception when lr_scheduler dict has unknown interval param value."""
model = BoringModel()
optimizer = optim.Adam(model.parameters())
model.configure_optimizers = lambda: {
"optimizer": optimizer,
"lr_scheduler": {"scheduler": optim.lr_scheduler.StepLR(optimizer, 1), "interval": "incorrect_unknown_value"},
}
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.raises(MisconfigurationException, match=r'The "interval" key in lr scheduler dict must be'):
trainer.fit(model)
def test_dataloader_source_request_from_module():
"""Test requesting a dataloader from a module works."""
module = BoringModel()
module.trainer = Trainer()
module.foo = Mock(return_value=module.train_dataloader())
source = _DataLoaderSource(module, "foo")
assert source.is_module()
module.foo.assert_not_called()
assert isinstance(source.dataloader(), DataLoader)
module.foo.assert_called_once()
def test_reducelronplateau_with_no_monitor_in_lr_scheduler_dict_raises(tmpdir):
"""Test exception when lr_scheduler dict has a ReduceLROnPlateau with no monitor."""
model = BoringModel()
optimizer = optim.Adam(model.parameters())
model.configure_optimizers = lambda: {
"optimizer": optimizer,
"lr_scheduler": {"scheduler": optim.lr_scheduler.ReduceLROnPlateau(optimizer)},
}
trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
with pytest.raises(MisconfigurationException, match="must include a monitor when a `ReduceLROnPlateau`"):
trainer.fit(model)
def test_trainer_manual_optimization_config(tmpdir):
"""Test error message when requesting Trainer features unsupported with manual optimization."""
model = BoringModel()
model.automatic_optimization = False
trainer = Trainer(gradient_clip_val=1.0)
with pytest.raises(MisconfigurationException, match="Automatic gradient clipping is not supported"):
trainer.fit(model)
trainer = Trainer(accumulate_grad_batches=2)
with pytest.raises(MisconfigurationException, match="Automatic gradient accumulation is not supported"):
trainer.fit(model)
