Decrease the number of epochs to 3 to prevent keras_premade_models_test from timing out with lengthy MultiWorkerMirroredStrategy tests. Otherwise it'd timeout with ~1 out of 100 runs.
@@ -107,14 +107,14 @@ class KerasPremadeModelsTest(tf.test.TestCase, parameterized.TestCase):
model.compile(opt, 'mse')
if use_dataset_creator:
x = dataset_creator.DatasetCreator(dataset_fn)
- hist = model.fit(x, epochs=5, steps_per_epoch=INPUT_SIZE)
+ hist = model.fit(x, epochs=3, steps_per_epoch=INPUT_SIZE)
else:
if data_fn == 'numpy':
inputs, output = get_numpy()
- hist = model.fit(inputs, output, epochs=5)
+ hist = model.fit(inputs, output, epochs=3)
else:
- hist = model.fit(get_dataset(), epochs=5)
- self.assertLess(hist.history['loss'][4], 0.2)
+ hist = model.fit(get_dataset(), epochs=3)
+ self.assertLess(hist.history['loss'][2], 0.2)
@tf.__internal__.distribute.combinations.generate(
strategy_combinations_eager_data_fn())
@@ -140,14 +140,14 @@ class KerasPremadeModelsTest(tf.test.TestCase, parameterized.TestCase):
if use_dataset_creator:
x = dataset_creator.DatasetCreator(dataset_fn)
- hist = wide_deep_model.fit(x, epochs=5, steps_per_epoch=INPUT_SIZE)
+ hist = wide_deep_model.fit(x, epochs=3, steps_per_epoch=INPUT_SIZE)
else:
if data_fn == 'numpy':
inputs, output = get_numpy()
- hist = wide_deep_model.fit(inputs, output, epochs=5)
+ hist = wide_deep_model.fit(inputs, output, epochs=3)
else:
- hist = wide_deep_model.fit(get_dataset(), epochs=5)
- self.assertLess(hist.history['loss'][4], 0.2)
+ hist = wide_deep_model.fit(get_dataset(), epochs=3)
+ self.assertLess(hist.history['loss'][2], 0.2)
if __name__ == '__main__':
