attempting to remove some speed issues (Lightning-AI#1482)

* removed some .items * added speed tests * added speed tests * Update benchmarks/test_rnn_parity.py Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com> * Update benchmarks/test_trainer_parity.py Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com> * fix lost model reference * added speed tests Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
tullie · May 6, 2020 · 47f07d6 · 47f07d6
1 parent 485d465
commit 47f07d6
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -63,6 +63,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added speed parity tests (max 1 sec difference per epoch)([#1482](https://github.com/PyTorchLightning/pytorch-lightning/pull/1482))
 - Added same step loggers' metrics aggregation ([#1278](https://github.com/PyTorchLightning/pytorch-lightning/pull/1278))
 - Added parity test between a vanilla MNIST model and lightning model ([#1284](https://github.com/PyTorchLightning/pytorch-lightning/pull/1284))
 - Added parity test between a vanilla RNN model and lightning model ([#1351](https://github.com/PyTorchLightning/pytorch-lightning/pull/1351))

diff --git a/benchmarks/test_rnn_parity.py b/benchmarks/test_rnn_parity.py
@@ -6,6 +6,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
+import tests.base.utils as tutils
 
 from pytorch_lightning import Trainer, LightningModule
 
@@ -64,6 +65,8 @@ def test_pytorch_parity(tmpdir):
     for pl_out, pt_out in zip(lightning_outs, manual_outs):
         np.testing.assert_almost_equal(pl_out, pt_out, 8)
 
+    tutils.assert_speed_parity(pl_times, pt_times, num_epochs)
+
 
 def set_seed(seed):
     np.random.seed(seed)

diff --git a/benchmarks/test_trainer_parity.py b/benchmarks/test_trainer_parity.py
@@ -8,6 +8,7 @@
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 from torchvision import transforms
+import tests.base.utils as tutils
 
 from pytorch_lightning import Trainer, LightningModule
 from tests.base.datasets import TestingMNIST
@@ -64,6 +65,8 @@ def test_pytorch_parity(tmpdir):
     for pl_out, pt_out in zip(lightning_outs, manual_outs):
         np.testing.assert_almost_equal(pl_out, pt_out, 5)
 
+    tutils.assert_speed_parity(pl_times, pt_times, num_epochs)
+
 
 def set_seed(seed):
     np.random.seed(seed)

diff --git a/pytorch_lightning/trainer/logging.py b/pytorch_lightning/trainer/logging.py
@@ -112,10 +112,6 @@ def process_output(self, output, train=False):
             num_gpus = self.num_gpus
             callback_metrics = self.reduce_distributed_output(callback_metrics, num_gpus)
 
-        for k, v in callback_metrics.items():
-            if isinstance(v, torch.Tensor):
-                callback_metrics[k] = v.item()
-
         # ---------------
         # EXTRACT PROGRESS BAR KEYS
         # ---------------

diff --git a/tests/base/utils.py b/tests/base/utils.py
@@ -20,6 +20,19 @@
 ROOT_PATH = os.path.abspath(os.path.dirname(__file__))
 
 
+def assert_speed_parity(pl_times, pt_times, num_epochs):
+
+    # assert speeds
+    max_diff_per_epoch = 0.9
+    pl_times = np.asarray(pl_times)
+    pt_times = np.asarray(pt_times)
+    diffs = pl_times - pt_times
+    diffs = diffs / num_epochs
+
+    assert np.alltrue(diffs < max_diff_per_epoch), \
+        f"lightning was slower than PT (threshold {max_diff_per_epoch})"
+
+
 def run_model_test_no_loggers(trainer_options, model, min_acc=0.50):
     # save_dir = trainer_options['default_root_dir']