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()