.10, #4

Author: Borda

benchmarks/test_rnn_parity.py

@@ -66,7 +66,7 @@ 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)
+    tutils.assert_speed_parity(pl_times, pt_times)
 
 
 def vanilla_loop(MODEL, num_runs=10, num_epochs=10):

benchmarks/test_trainer_parity.py

@@ -52,7 +52,7 @@ def test_pytorch_parity(tmpdir):
     :param tmpdir:
     :return:
     """
-    num_epochs = 5
+    num_epochs = 4
     num_rums = 3
     lightning_outs, pl_times = lightning_loop(ParityModuleMNIST, num_rums, num_epochs)
     manual_outs, pt_times = vanilla_loop(ParityModuleMNIST, num_rums, num_epochs)
@@ -62,7 +62,7 @@ def test_pytorch_parity(tmpdir):
         np.testing.assert_almost_equal(pl_out, pt_out, 5)
 
     # the fist run initialize dataset (download & filter)
-    tutils.assert_speed_parity(pl_times[1:], pt_times[1:], num_epochs)
+    tutils.assert_speed_parity(pl_times[1:], pt_times[1:])
 
 
 def vanilla_loop(cls_model, num_runs=10, num_epochs=10):

tests/base/utils.py

@@ -12,11 +12,11 @@
 from tests.base.model_template import EvalModelTemplate
 
 
-def assert_speed_parity(pl_times, pt_times, num_epochs, max_diff_per_epoch=0.05):
+def assert_speed_parity(pl_times, pt_times, max_diff_per_epoch=0.1):
     # assert speeds
     diffs = np.asarray(pl_times) - np.asarray(pt_times)
-    # norm by nb epochs and tha vanila time
-    diffs = diffs / num_epochs / np.asarray(pt_times)
+    # norm by vanila time
+    diffs = diffs / np.asarray(pt_times)
     assert np.alltrue(diffs < max_diff_per_epoch), \
         f"lightning {diffs} was slower than PT (threshold {max_diff_per_epoch})"