feat(trainer): add enable_benchmarking option (#803)

* feat(trainer): add enable_benchmarking option

closes #370

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* add test

* try to make the lint work

* fix typo

* add test, verify torch.backends.cudnn.benchmark

* make lint happy

* make lint happy

Co-authored-by: William Falcon <waf2107@columbia.edu>

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]
+                <https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936>`_.
+
         .. 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:

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