Skip to content

Commit

Permalink
[Test] Speed up lookahead_test (#1169)
Browse files Browse the repository at this point in the history
Reduce number of samples from 10e5 to 10e4, increase error from 1e-4 to 1e-3

CPU test time: 33s -> 17s
  • Loading branch information
Squadrick authored Feb 27, 2020
1 parent 83e10cb commit 487e366
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions tensorflow_addons/optimizers/lookahead_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,13 +109,13 @@ def test_fit_simple_linear_model(self):
np.random.seed(0x2019)
tf.random.set_seed(0x2019)

x = np.random.standard_normal((100000, 3))
x = np.random.standard_normal((10000, 3))
w = np.random.standard_normal((3, 1))
y = np.dot(x, w) + np.random.standard_normal((100000, 1)) * 1e-4
y = np.dot(x, w) + np.random.standard_normal((10000, 1)) * 1e-4

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
model.compile(Lookahead("adam"), loss="mse")
model.compile(Lookahead("sgd"), loss="mse")

model.fit(x, y, epochs=3)

Expand All @@ -124,7 +124,7 @@ def test_fit_simple_linear_model(self):
predicted = model.predict(x)

max_abs_diff = np.max(np.abs(predicted - y))
self.assertLess(max_abs_diff, 1e-4)
self.assertLess(max_abs_diff, 1e-3)

def test_model_dynamic_lr(self):
grad = tf.Variable([[0.1]])
Expand Down

0 comments on commit 487e366

Please sign in to comment.