tests: add test for tf keras Native parallel training
@@ -85,6 +85,26 @@ class NativeImplementations:
min_num_gpus_required=0,
)
+ # Train a single tf.keras model using fit() on multiple GPUs.
+ TFKerasMNISTCNNSingleFitParallel = NativeImplementation(
+ cwd=conf.experimental_path("native_mnist_tf_keras"),
+ command=[
+ "python",
+ conf.experimental_path("native_mnist_tf_keras/native_impl.py"),
+ "--use-fit",
+ ],
+ configuration={
+ "batches_per_step": 4,
+ "checkpoint_storage": experiment.shared_fs_checkpoint_config(),
+ "searcher": {"name": "single", "max_steps": 1, "metric": "val_accuracy"},
+ "resources": {"slots_per_trial": 2},
+ "max_restarts": 0,
+ },
+ num_expected_steps_per_trial=1,
+ num_expected_trials=1,
+ min_num_gpus_required=2,
+ )
+
# Train a single tf.keras model using fit_generator().
TFKerasMNISTCNNSingleFitGenerator = NativeImplementation(
cwd=conf.experimental_path("native_mnist_tf_keras"),
@@ -199,6 +219,7 @@ def run_test_mode(implementation: NativeImplementation) -> None:
NativeImplementations.TFKerasMNISTCNNSingleFitGenerator,
NativeImplementations.TFKerasMNISTCNNSingleFit,
NativeImplementations.TFKerasMNISTCNNSingleGeneric,
+ NativeImplementations.TFKerasMNISTCNNSingleFitParallel,
],
)
@pytest.mark.parametrize("tf2", [True, False]) # type: ignore
