From 460a102b61c831d1814aec0dc0ce4c6b65722377 Mon Sep 17 00:00:00 2001
From: gabrieldemarmiesse <gabrieldemarmiesse@gmail.com>
Date: Sun, 23 Feb 2020 18:52:31 +0000
Subject: [PATCH] Made the rectified adam test faster.

---
 .../optimizers/rectified_adam_test.py         | 57 +++++++++----------
 1 file changed, 27 insertions(+), 30 deletions(-)

diff --git a/tensorflow_addons/optimizers/rectified_adam_test.py b/tensorflow_addons/optimizers/rectified_adam_test.py
index f646e798d8..8471e4e15f 100644
--- a/tensorflow_addons/optimizers/rectified_adam_test.py
+++ b/tensorflow_addons/optimizers/rectified_adam_test.py
@@ -69,19 +69,18 @@ def run_sparse_sample(self, iterations, expected, optimizer):
         self.assertAllClose(var_1.read_value(), expected[1], atol=2e-4)
 
     def test_dense_sample(self):
-        # Expected values are obtained from the official implementation
+        # Expected values are obtained from the previous implementation
         self.run_dense_sample(
-            iterations=1000,
-            expected=[[0.5554, 1.5549], [2.5557, 3.5557]],
+            iterations=100,
+            expected=[[0.985769, 1.985269], [2.986119, 3.986068]],
             optimizer=RectifiedAdam(lr=1e-3),
         )
 
     def test_sparse_sample(self):
-        # Expected values are obtained from the official implementation
-        # Dense results should be: [-0.1929,  0.8066], [1.8075, 2.8074]
+        # Expected values are obtained from the previous implementation
         self.run_sparse_sample(
-            iterations=2000,
-            expected=[[-0.1929, 2.0], [3.0, 2.8074]],
+            iterations=200,
+            expected=[[0.959333, 2.0], [3.0, 3.959632]],
             optimizer=RectifiedAdam(lr=1e-3),
         )
 
@@ -89,8 +88,8 @@ def test_dense_sample_with_amsgrad(self):
         # Expected values are obtained from the official implementation
         # `amsgrad` has no effect because the gradient is fixed
         self.run_dense_sample(
-            iterations=1000,
-            expected=[[0.5554, 1.5549], [2.5557, 3.5557]],
+            iterations=100,
+            expected=[[0.985769, 1.985269], [2.986119, 3.986068]],
             optimizer=RectifiedAdam(lr=1e-3, amsgrad=True),
         )
 
@@ -98,43 +97,42 @@ def test_sparse_sample_with_amsgrad(self):
         # Expected values are obtained from the official implementation
         # `amsgrad` has no effect because the gradient is fixed
         self.run_sparse_sample(
-            iterations=2000,
-            expected=[[-0.1929, 2.0], [3.0, 2.8074]],
+            iterations=200,
+            expected=[[0.959333, 2.0], [3.0, 3.959632]],
             optimizer=RectifiedAdam(lr=1e-3, amsgrad=True),
         )
 
     def test_dense_sample_with_weight_decay(self):
-        # Expected values are obtained from the official implementation
+        # Expected values are obtained from the previous implementation
         self.run_dense_sample(
-            iterations=1000,
-            expected=[[0.5472, 1.5368], [2.5276, 3.5176]],
+            iterations=100,
+            expected=[[0.984775, 1.983276], [2.983125, 3.982076]],
             optimizer=RectifiedAdam(lr=1e-3, weight_decay=0.01),
         )
 
     def test_sparse_sample_with_weight_decay(self):
-        # Expected values are obtained from the official implementation
-        # Dense results should be: [-0.2029,  0.7768], [1.7578, 2.7380]
+        # Expected values are obtained from the previous implementation
         self.run_sparse_sample(
-            iterations=2000,
-            expected=[[-0.2029, 2.0], [3.0, 2.7380]],
+            iterations=200,
+            expected=[[0.957368, 2.0], [3.0, 3.951673]],
             optimizer=RectifiedAdam(lr=1e-3, weight_decay=0.01),
         )
 
     def test_dense_sample_with_warmup(self):
         self.run_dense_sample(
-            iterations=1000,
-            expected=[[0.8041, 1.8041], [2.8041, 3.8041]],
+            iterations=100,
+            expected=[[0.994062, 1.993912], [2.994167, 3.994152]],
             optimizer=RectifiedAdam(
-                lr=1e-3, total_steps=1000, warmup_proportion=0.1, min_lr=1e-5,
+                lr=1e-3, total_steps=100, warmup_proportion=0.1, min_lr=1e-5,
             ),
         )
 
     def test_sparse_sample_with_warmup(self):
         self.run_sparse_sample(
-            iterations=2000,
-            expected=[[0.4653, 2.0], [3.0, 3.4653]],
+            iterations=200,
+            expected=[[0.982629, 2.0], [3.0, 3.982674]],
             optimizer=RectifiedAdam(
-                lr=1e-3, total_steps=2000, warmup_proportion=0.1, min_lr=1e-5,
+                lr=1e-3, total_steps=200, warmup_proportion=0.1, min_lr=1e-5,
             ),
         )
 
@@ -142,8 +140,8 @@ def test_dense_sample_with_lookahead(self):
         # Expected values are obtained from the original implementation
         # of Ranger
         self.run_dense_sample(
-            iterations=1000,
-            expected=[[0.7985, 1.7983], [2.7987, 3.7986]],
+            iterations=100,
+            expected=[[0.993126, 1.992901], [2.993283, 3.993261]],
             optimizer=Lookahead(
                 RectifiedAdam(lr=1e-3, beta_1=0.95,),
                 sync_period=6,
@@ -152,12 +150,11 @@ def test_dense_sample_with_lookahead(self):
         )
 
     def test_sparse_sample_with_lookahead(self):
-        # Expected values are obtained from the original implementation
+        # Expected values are obtained from the previous implementation
         # of Ranger.
-        # Dense results should be: [0.6417,  1.6415], [2.6419, 3.6418]
         self.run_sparse_sample(
-            iterations=1500,
-            expected=[[0.6417, 2.0], [3.0, 3.6418]],
+            iterations=150,
+            expected=[[0.988156, 2.0], [3.0, 3.988291]],
             optimizer=Lookahead(
                 RectifiedAdam(lr=1e-3, beta_1=0.95,),
                 sync_period=6,