Made the rectified adam test faster. (#1142)

Author: gabrieldemarmiesse

tensorflow_addons/optimizers/rectified_adam_test.py

@@ -69,81 +69,79 @@ 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),
         )
 
     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),
         )
 
     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,
             ),
         )
 
     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,