Keras file format updates (#1401)

* omlt import fix

* updated main surrogate code; baseline tests fixed with re-run NNs

* black and filename option

* remove version restriction on tensorflow

* removing unused vars

* testing model format fixes

* updating additional tests

* rerun black

* setup.py fixes

* black on setup.py

* expected test failures

* rerun black

* correct version of black

* deprecation warning

* update deprecation status

* skip testing on keras_surrogates for python 3.8

* black

* remove old keras files

* remove more old keras files

* try removing xfails from keras surrogate tests

* add comments explaining xfailing sofc surrogate tests

* Try removing workaround for TensorFlow failures for Python 3.11

* Add .keras suffix to package_data

---------

Co-authored-by: Brandon Paul <[email protected]>
Co-authored-by: Andrew Lee <[email protected]>
Co-authored-by: Ludovico Bianchi <[email protected]>

.github/actions/run-examples/examples_for_idaes_ci.py

@@ -28,24 +28,13 @@ def _matches_pattern(item: pytest.Item, pattern: str) -> bool:
 
 
 def pytest_configure(config: pytest.Config):
-    tensorflow_py311_win = pytest.mark.xfail(
-        condition=sys.version_info > (3, 11),
-        run=True,
-        strict=False,
-        reason="tensorflow ImportError on 3.11+ on Windows (cannot import name 'formatargspec' from 'inspect')",
-    )
     config.stash[matchmarker] = {
         "*/held/*": pytest.mark.xfail(run=False, reason="notebook has 'held' status"),
         "*/archive/*": pytest.mark.skip(reason="notebook is archived"),
-        # TODO: Need to fix this once the Python 3.11 issue is resolved in tensorflow
-        "*/surrogates/best_practices_optimization*": tensorflow_py311_win,
-        "*/surrogates/omlt/keras_flowsheet_optimization*": tensorflow_py311_win,
         "*/surrogates/sco2/alamo/*": pytest.mark.xfail(
             run=False,
             reason="notebooks require ALAMO to run",
         ),
-        "*/surrogates/sco2/omlt/keras_training*": tensorflow_py311_win,
-        "*/surrogates/sco2/omlt/flowsheet_optimization*": tensorflow_py311_win,
     }
     config.stash[marked] = []
 

idaes/core/surrogate/keras_surrogate.py

@@ -34,7 +34,7 @@
 
 if omlt_available:
     from omlt import OmltBlock, OffsetScaling
-    from omlt.neuralnet import (
+    from omlt.neuralnet.nn_formulation import (
         FullSpaceSmoothNNFormulation,
         ReducedSpaceSmoothNNFormulation,
         ReluBigMFormulation,
@@ -253,7 +253,7 @@ def evaluate_surrogate(self, inputs):
             y = self._output_scaler.unscale(y)
         return y
 
-    def save_to_folder(self, keras_folder_name):
+    def save_to_folder(self, keras_folder_name, keras_model_name="idaes_keras_model"):
         """
         Save the surrogate object to disk by providing the name of the
         folder to contain the keras model and additional IDAES metadata
@@ -263,7 +263,9 @@ def save_to_folder(self, keras_folder_name):
               The name of the folder to contain the Keras model and additional
               IDAES metadata
         """
-        self._keras_model.save(keras_folder_name)
+        self._keras_model.save(
+            os.path.join(keras_folder_name, keras_model_name + ".keras")
+        )
         info = dict()
         info["input_scaler"] = None
         if self._input_scaler is not None:
@@ -281,7 +283,7 @@ def save_to_folder(self, keras_folder_name):
             json.dump(info, fd)
 
     @classmethod
-    def load_from_folder(cls, keras_folder_name):
+    def load_from_folder(cls, keras_folder_name, keras_model_name="idaes_keras_model"):
         """
         Load the surrogate object from disk by providing the name of the
         folder holding the keras model
@@ -293,7 +295,11 @@ def load_from_folder(cls, keras_folder_name):
 
         Returns: an instance of KerasSurrogate
         """
-        keras_model = keras.models.load_model(keras_folder_name)
+
+        keras_model = keras.models.load_model(
+            os.path.join(keras_folder_name, keras_model_name + ".keras")
+        )
+
         with open(os.path.join(keras_folder_name, "idaes_info.json")) as fd:
             info = json.load(fd)
 
@@ -319,12 +325,11 @@ def save_keras_json_hd5(nn, path, name):
     json_model = nn.to_json()
     with open(os.path.join(path, "{}.json".format(name)), "w") as json_file:
         json_file.write(json_model)
-    nn.save_weights(os.path.join(path, "{}.h5".format(name)))
+    nn.save(os.path.join(path, "{}.keras".format(name)))
+    nn.save_weights(os.path.join(path, "{}.weights.h5".format(name)))
 
 
 def load_keras_json_hd5(path, name):
-    with open(os.path.join(path, "{}.json".format(name)), "r") as json_file:
-        json_model = json_file.read()
-        nn = keras.models.model_from_json(json_model)
-    nn.load_weights(os.path.join(path, "{}.h5".format(name)))
+    nn = keras.models.load_model(os.path.join(path, "{}.keras".format(name)))
+    nn.load_weights(os.path.join(path, "{}.weights.h5".format(name)))
     return nn

idaes/core/surrogate/tests/data/create_keras_models.py

@@ -82,7 +82,7 @@ def compare_fit(nn, x_test, y_test, input_scaler, output_scaler):
 nn.compile(optimizer=Adam(), loss="mse")
 
 mcp_save = ModelCheckpoint(
-    ".mdl_wts.hdf5", save_best_only=True, monitor="val_loss", mode="min"
+    ".mdl_wts.keras", save_best_only=True, monitor="val_loss", mode="min"
 )
 history = nn.fit(
     x=x,
@@ -111,7 +111,7 @@ def compare_fit(nn, x_test, y_test, input_scaler, output_scaler):
 nn.compile(optimizer=Adam(), loss="mse")
 
 mcp_save = ModelCheckpoint(
-    ".mdl_wts.hdf5", save_best_only=True, monitor="val_loss", mode="min"
+    ".mdl_wts.keras", save_best_only=True, monitor="val_loss", mode="min"
 )
 history = nn.fit(
     x=x,
@@ -161,7 +161,7 @@ def compare_fit(nn, x_test, y_test, input_scaler, output_scaler):
 nn.compile(optimizer=Adam(), loss="mse")
 
 mcp_save = ModelCheckpoint(
-    ".mdl_wts.hdf5", save_best_only=True, monitor="val_loss", mode="min"
+    ".mdl_wts.keras", save_best_only=True, monitor="val_loss", mode="min"
 )
 history = nn.fit(
     x=x,
@@ -188,7 +188,7 @@ def compare_fit(nn, x_test, y_test, input_scaler, output_scaler):
 nn.compile(optimizer=Adam(), loss="mse")
 
 mcp_save = ModelCheckpoint(
-    ".mdl_wts.hdf5", save_best_only=True, monitor="val_loss", mode="min"
+    ".mdl_wts.keras", save_best_only=True, monitor="val_loss", mode="min"
 )
 history = nn.fit(
     x=x,

idaes/core/surrogate/tests/data/keras_models/PT_data_2_10_10_2_relu.json

@@ -1 +1 @@
-{"class_name": "Sequential", "config": {"name": "PT_data_2_10_10_2_relu", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": [null, 2], "dtype": "float32", "sparse": false, "ragged": false, "name": "dense_9_input"}}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": true, "batch_input_shape": [null, 2], "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": true, "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": true, "dtype": "float32", "units": 2, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}]}, "keras_version": "2.7.0", "backend": "tensorflow"}
+{"module": "keras", "class_name": "Sequential", "config": {"name": "PT_data_2_10_10_2_relu", "trainable": true, "dtype": "float32", "layers": [{"module": "keras.layers", "class_name": "InputLayer", "config": {"batch_shape": [null, 2], "dtype": "float32", "sparse": false, "name": "input_layer_3"}, "registered_name": null}, {"module": "keras.layers", "class_name": "Dense", "config": {"name": "dense_9", "trainable": true, "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"module": "keras.initializers", "class_name": "GlorotUniform", "config": {"seed": null}, "registered_name": null}, "bias_initializer": {"module": "keras.initializers", "class_name": "Zeros", "config": {}, "registered_name": null}, "kernel_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "registered_name": null, "build_config": {"input_shape": [null, 2]}}, {"module": "keras.layers", "class_name": "Dense", "config": {"name": "dense_10", "trainable": true, "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"module": "keras.initializers", "class_name": "GlorotUniform", "config": {"seed": null}, "registered_name": null}, "bias_initializer": {"module": "keras.initializers", "class_name": "Zeros", "config": {}, "registered_name": null}, "kernel_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "registered_name": null, "build_config": {"input_shape": [null, 10]}}, {"module": "keras.layers", "class_name": "Dense", "config": {"name": "dense_11", "trainable": true, "dtype": "float32", "units": 2, "activation": "linear", "use_bias": true, "kernel_initializer": {"module": "keras.initializers", "class_name": "GlorotUniform", "config": {"seed": null}, "registered_name": null}, "bias_initializer": {"module": "keras.initializers", "class_name": "Zeros", "config": {}, "registered_name": null}, "kernel_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "registered_name": null, "build_config": {"input_shape": [null, 10]}}], "build_input_shape": [null, 2]}, "registered_name": null, "build_config": {"input_shape": [null, 2]}, "compile_config": {"optimizer": {"module": "keras.optimizers", "class_name": "Adam", "config": {"name": "adam", "learning_rate": 0.0010000000474974513, "weight_decay": null, "clipnorm": null, "global_clipnorm": null, "clipvalue": null, "use_ema": false, "ema_momentum": 0.99, "ema_overwrite_frequency": null, "loss_scale_factor": null, "gradient_accumulation_steps": null, "beta_1": 0.9, "beta_2": 0.999, "epsilon": 1e-07, "amsgrad": false}, "registered_name": null}, "loss": "mse", "loss_weights": null, "metrics": null, "weighted_metrics": null, "run_eagerly": false, "steps_per_execution": 1, "jit_compile": false}}