diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 815f2d9d..566c2e73 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -110,6 +110,7 @@ class Trainer(TrainerIOMixin, truncated_bptt_steps: Optional[int] = None, resume_from_checkpoint: Optional[str] = None, profiler: Optional[BaseProfiler] = None, + benchmark: bool = False, reload_dataloaders_every_epoch: bool = False, ): r""" @@ -583,12 +584,26 @@ class Trainer(TrainerIOMixin, trainer = Trainer(profiler=profiler) reload_dataloaders_every_epoch: Set to True to reload dataloaders every epoch + benchmark (bool): If true enables cudnn.benchmark. + This flag is likely to increase the speed of your system if your + input sizes don't change. However, if it does, then it will likely + make your system slower. + + The speedup comes from allowing the cudnn auto-tuner to find the best + algorithm for the hardware `[see discussion here] + `_. + .. warning:: Following arguments become deprecated and they will be removed in v0.8.0: - `nb_sanity_val_steps` """ + # benchmarking + self.benchmark = benchmark + if benchmark: + torch.backends.cudnn.benchmark = True + # Transfer params # Backward compatibility if nb_gpu_nodes is not None: diff --git a/tests/test_trainer.py b/tests/test_trainer.py index e5d7f68c..f33b1f3e 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -793,6 +793,40 @@ def test_trainer_min_steps_and_epochs(tmpdir): trainer.current_epoch > 0, "Model did not train for at least min_steps" +def test_benchmark_option(tmpdir): + """Verify benchmark option.""" + tutils.reset_seed() + + class CurrentTestModel( + LightningValidationMultipleDataloadersMixin, + LightningTestModelBase + ): + pass + + hparams = tutils.get_hparams() + model = CurrentTestModel(hparams) + + # verify torch.backends.cudnn.benchmark is not turned on + assert not torch.backends.cudnn.benchmark + + # logger file to get meta + trainer_options = dict( + default_save_path=tmpdir, + max_epochs=1, + benchmark=True, + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # verify training completed + assert result == 1 + + # verify torch.backends.cudnn.benchmark is not turned off + assert torch.backends.cudnn.benchmark + + def test_testpass_overrides(tmpdir): hparams = tutils.get_hparams()