From 600f8e4c818d99bd937798ed55c42133fe656056 Mon Sep 17 00:00:00 2001 From: Szubie Date: Fri, 15 Oct 2021 20:31:50 +0100 Subject: [PATCH] Add a test for TF SavedModel format saving and loading --- tests/test_model_saving.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_model_saving.py b/tests/test_model_saving.py index 692ebeb..ad34ba8 100644 --- a/tests/test_model_saving.py +++ b/tests/test_model_saving.py @@ -130,8 +130,8 @@ def _supervised_custom_model_saving(model_filepath, save_fn, load_fn): y_pred_2 = model_2.fit_transform(X, Y) ### Save and load ### -def _save_ivis_model(model, filepath): - model.save_model(filepath, overwrite=True) +def _save_ivis_model(model, filepath, save_format='h5'): + model.save_model(filepath, save_format=save_format, overwrite=True) def _load_ivis_model(filepath): model_2 = Ivis() @@ -179,6 +179,11 @@ def _undill_ivis_model(filepath): test_supervised_custom_model_pickling = partial(_supervised_custom_model_saving, save_fn=_dill_ivis_model, load_fn=_undill_ivis_model) +### Other ### +test_tf_savedmodel_persistence = partial(_unsupervised_model_save_test, + save_fn=partial(_save_ivis_model, save_format='tfs'), + load_fn=_load_ivis_model) + def test_save_overwriting(model_filepath): model = Ivis(k=15, batch_size=16, epochs=2) iris = datasets.load_iris()