Call base_model.save() for backward compatibility.

csferng · csferng · commit 47d66e23a35a · 2022-06-09T13:07:52.000-07:00
The save() interface has changed in different TF versions. Since NSL
supports back to TF 1.15, it is better to accept variadic arguments.

Also updated the example notebook following internal review comments.
diff --git a/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb b/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb
@@ -6,7 +6,7 @@
         "id": "ZwZNOAMZcxl3"
       },
       "source": [
-        "##### Copyright 2019 The TensorFlow Neural Structured Learning Authors"
+        "##### Copyright 2022 The TensorFlow Neural Structured Learning Authors"
       ]
     },
     {
@@ -37,7 +37,7 @@
         "id": "-BszoQj0dSZO"
       },
       "source": [
-        "# How to save and load models in NSL\n"
+        "# Saving and loading NSL models\n"
       ]
     },
     {
@@ -48,10 +48,10 @@
       "source": [
         "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
         "  <td>\n",
-        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/How_to_save_and_load_model_in_nsl.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
+        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/How_to_save_and_load_model_in_nsl.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
+        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
         "  </td>\n",
         "</table>"
       ]
@@ -76,14 +76,14 @@
         "    * Here we create the base model with the `tf.keras` functional API.\n",
         "    * This procedure is also compatible with models created by the\n",
         "      `tf.keras.Sequential` API and models created by subclasses of\n",
-        "      `tf.keras.Model`, provided that the subclasses support saving and loading.\n",
+        "      `tf.keras.Model`, provided that the subclass models support saving and loading.\n",
         "    * For more information on Keras models in TensorFlow, see this\n",
         "      [documentation](https://www.tensorflow.org/api_docs/python/tf/keras/Model).\n",
         "2.  Wrap the base model with the **`AdversarialRegularization`** wrapper class.\n",
         "    * The wrapper class is provided by the NSL framework, and it creates a new\n",
         "      `tf.keras.Model` instance.\n",
-        "    * This new model will include the adversarial loss as a regularization term\n",
-        "      in its training objective.\n",
+        "    * This new model is a subclass model and will include the adversarial loss\n",
+        "      as a regularization term in its training objective.\n",
         "3.  Train and evaluate the new model.\n",
         "4.  Save the model, and load it back.\n",
         "5.  Verify that the loaded model is the same as the one before saving."
@@ -107,17 +107,6 @@
         "Install the Neural Structured Learning package."
       ]
     },
-    {
-      "cell_type": "code",
-      "source": [
-        "!pip uninstall -y neural-structured-learning"
-      ],
-      "metadata": {
-        "id": "rYfawinCPpuC"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -135,7 +124,7 @@
         "id": "PZvsEQrhSqKx"
       },
       "source": [
-        "Import libraries. We abbreviate `neural_structured_learning` to `nsl`."
+        "Import libraries."
       ]
     },
     {
@@ -683,7 +672,7 @@
     "accelerator": "GPU",
     "colab": {
       "collapsed_sections": [],
-      "name": "How to save and load model in nsl.ipynb",
+      "name": "Saving and loading NSL models",
       "private_outputs": true,
       "provenance": [],
       "toc_visible": true
diff --git a/neural_structured_learning/keras/adversarial_regularization.py b/neural_structured_learning/keras/adversarial_regularization.py
@@ -23,10 +23,6 @@
 
 
 import attr
-
-from keras.saving import save
-from keras.utils import traceback_utils
-
 import neural_structured_learning.configs as nsl_configs
 from neural_structured_learning.lib import adversarial_neighbor
 import six
@@ -699,58 +695,11 @@ def call(self, inputs, **kwargs):
         scaled_adv_loss, name='scaled_adversarial_loss', aggregation='mean')
     return outputs
 
-
-  @traceback_utils.filter_traceback
-  def save(self,
-           filepath,
-           overwrite=True,
-           include_optimizer=True,
-           save_format=None,
-           signatures=None,
-           options=None,
-           save_traces=True):
-    # pylint: disable=line-too-long
-    """Saves the model to Tensorflow SavedModel or a single HDF5 file.
-    Please see `tf.keras.models.save_model` or the
-    [Serialization and Saving guide](https://keras.io/guides/serialization_and_saving/)
-    for details.
-    Args:
-        filepath: String, PathLike, path to SavedModel or H5 file to save the
-            model.
-        overwrite: Whether to silently overwrite any existing file at the
-            target location, or provide the user with a manual prompt.
-        include_optimizer: If True, save optimizer's state together.
-        save_format: Either `'tf'` or `'h5'`, indicating whether to save the
-            model to Tensorflow SavedModel or HDF5. Defaults to 'tf' in TF 2.X,
-            and 'h5' in TF 1.X.
-        signatures: Signatures to save with the SavedModel. Applicable to the
-            'tf' format only. Please see the `signatures` argument in
-            `tf.saved_model.save` for details.
-        options: (only applies to SavedModel format)
-            `tf.saved_model.SaveOptions` object that specifies options for
-            saving to SavedModel.
-        save_traces: (only applies to SavedModel format) When enabled, the
-            SavedModel will store the function traces for each layer. This
-            can be disabled, so that only the configs of each layer are stored.
-            Defaults to `True`. Disabling this will decrease serialization time
-            and reduce file size, but it requires that all custom layers/models
-            implement a `get_config()` method.
-    Example:
-    ```python
-    from keras.models import load_model
-    model.save('my_model.h5')  # creates a HDF5 file 'my_model.h5'
-    del model  # deletes the existing model
-    # returns a compiled model
-    # identical to the previous one
-    model = load_model('my_model.h5')
-    ```
-    """
-    # pylint: enable=line-too-long
-    save.save_model(self.base_model, filepath, overwrite, include_optimizer, save_format,
-                    signatures, options, save_traces)
-
-
-
+  def save(self, *args, **kwargs):
+    """Saves the base model. See base class for details of the interface."""
+    # Adversarial regularization doesn't introduce new model variables, so
+    # saving the base model can capture all variables in the model.
+    self.base_model.save(*args, **kwargs)
 
   def perturb_on_batch(self, x, **config_kwargs):
     """Perturbs the given input to generates adversarial examples.
diff --git a/neural_structured_learning/requirements.txt b/neural_structured_learning/requirements.txt
@@ -2,4 +2,3 @@ absl-py>=0.7
 attrs>=19.1.0
 scipy>=1.2.0
 six>=1.10
-keras>= 2.2.5
