From b78faedf5eb5aeb63450c13dec437f79811983a8 Mon Sep 17 00:00:00 2001 From: Paul Chiang Date: Thu, 21 Oct 2021 12:32:37 -0700 Subject: [PATCH] Use isdir() to decide whether to load_model as hdf5 or SavedModel PiperOrigin-RevId: 404855256 --- keras/saving/save.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/keras/saving/save.py b/keras/saving/save.py index 55915f0785d..cf40893f220 100644 --- a/keras/saving/save.py +++ b/keras/saving/save.py @@ -198,21 +198,24 @@ def load_model(filepath, custom_objects=None, compile=True, options=None): # py with generic_utils.SharedObjectLoadingScope(): with generic_utils.CustomObjectScope(custom_objects or {}): with load_context.load_context(options): - if (h5py is not None and - (isinstance(filepath, h5py.File) or h5py.is_hdf5(filepath))): + filepath_str = path_to_string(filepath) + if isinstance(filepath_str, str): + if not tf.io.gfile.exists(filepath_str): + raise IOError(f'No file or directory found at {filepath_str}') + + if tf.io.gfile.isdir(filepath_str): + return saved_model_load.load(filepath_str, compile, options) + else: + if h5py is None: + raise ImportError( + 'Filepath looks like a hdf5 file but h5py is not available.' + f' filepath={filepath}') + return hdf5_format.load_model_from_hdf5(filepath, custom_objects, + compile) + elif h5py is not None and isinstance(filepath, h5py.File): return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile) - filepath = path_to_string(filepath) - if isinstance(filepath, str): - if not tf.io.gfile.exists(filepath): - raise IOError(f'No file or directory found at {filepath}') - if saving_utils.is_hdf5_filepath(filepath) and h5py is None: - raise ImportError( - 'Filepath looks like a hdf5 file but h5py is not available.' - f' filepath={filepath}') - return saved_model_load.load(filepath, compile, options) - raise IOError( 'Unable to load model. Filepath is not an hdf5 file (or h5py is not ' f'available) or SavedModel. Received: filepath={filepath}')