def test_test_epoch_end(self, model: pl.LightningModule):
r"""Tests that ``test_epoch_end()`` runs and outputs a dict as required by PyTorch
Lightning.
Because of the size of some models, this test is only run when a GPU is available.
"""
if isinstance(model, LightningDistributedDataParallel):
pytest.skip()
check_overriden(model, "test_dataloader")
check_overriden(model, "test_step")
check_overriden(model, "test_epoch_end")
dl = model.test_dataloader()
if torch.cuda.is_available():
model = model.cuda()
model.eval()
outputs = [
model.test_step([x.cuda() for x in batch], 0) for batch in dl
]
else:
model.eval()
outputs = [model.test_step(batch, 0) for batch in dl]
result = model.test_epoch_end(outputs)
assert isinstance(result, dict)
return outputs, result
def test_forward(self, model: pl.LightningModule, data: torch.Tensor, training: bool):
r"""Calls ``model.forward()`` and tests that the output is not ``None``.
Because of the size of some models, this test is only run when a GPU is available.
"""
if isinstance(model, (LightningDistributedDataParallel, LightningDistributedModule)):
pytest.skip()
if torch.cuda.is_available():
model = model.cuda() # type: ignore
data = data.cuda()
if training:
model.train()
else:
model.eval()
_ = model(data)
assert _ is not None
def test_validation_step(self, model: pl.LightningModule):
r"""Runs a validation step based on the data returned from ``model.val_dataloader()``.
Tests that the dictionary returned from ``validation_step()`` are as required by PyTorch
Lightning.
Because of the size of some models, this test is only run when a GPU is available.
"""
if isinstance(model, LightningDistributedDataParallel):
pytest.skip()
check_overriden(model, "val_dataloader")
check_overriden(model, "validation_step")
dl = model.val_dataloader()
# TODO this can't handle multiple optimizers
batch = next(iter(dl))
if torch.cuda.is_available():
batch = [x.cuda() for x in batch]
model = model.cuda()
model.eval()
output = model.validation_step(batch, 0)
assert isinstance(output, dict)
if "loss" in output.keys():
assert isinstance(output["loss"], Tensor)
assert output["loss"].shape == (1, )
if "log" in output.keys():
assert isinstance(output["log"], dict)
if "progress_bar" in output.keys():
assert isinstance(output["progress_bar"], dict)
return batch, output
def test_training_step(self, model: pl.LightningModule, distributed: bool):
r"""Runs a training step based on the data returned from ``model.train_dataloader()``.
Tests that the dictionary returned from ``training_step()`` are as required by PyTorch
Lightning. A backward pass and optimizer step are also performed using the optimizer
provided by :func:`LightningModule.configure_optimizers`. By default, training steps
are tested for distributed and non-distributed models using the
:class:`torch.nn.parallel.DistributedDataParallel` wrapper. Distributed tests can be disabled
by setting :attr:`LightningModuleTest.DISTRIBUTED` to ``False``.
Because of the size of some models, this test is only run when a GPU is available.
"""
if distributed:
if not self.DISTRIBUTED:
pytest.skip(
"LightningModuleTest.DISTRIBUTED was False, skipping distributed training step"
)
self.init_process_group()
model = self._distributed_model(model.cuda())
if isinstance(model, LightningDistributedDataParallel):
dl = model.module.train_dataloader()
else:
dl = model.train_dataloader()
batch = next(iter(dl))
if torch.cuda.is_available():
batch = [x.cuda() for x in batch]
model = model.cuda()
model.train()
# TODO this can't handle multiple optimizers
if isinstance(model, LightningDistributedDataParallel):
output = model(batch, 0)
else:
output = model.training_step(batch, 0)
assert isinstance(output, dict)
assert "loss" in output.keys(), "loss key is required"
assert_valid_loss(output["loss"])
# test loss / backward pass, important for distributed operation
if isinstance(model, LightningDistributedDataParallel):
output = model(batch, 0)
optimizers, lr_scheduler, frequencies = _ProcessOptimizers(
).init_optimizers(model.module)
else:
optimizers, lr_scheduler, frequencies = _ProcessOptimizers(
).init_optimizers(model)
optim = optimizers[0]
loss = output["loss"]
optim.zero_grad()
loss.backward()
optim.step()
if "log" in output.keys():
assert isinstance(output["log"], dict)
if "progress_bar" in output.keys():
assert isinstance(output["progress_bar"], dict)
return batch, output
def model_to_device(self, model: LightningModule):
"""Moves the model to the appropriate device."""
if self.use_gpu:
model.cuda(self.trainer.root_gpu)
else:
model.cpu()