Update adversarial_regularization.py

Author: wangbingnan136

neural_structured_learning/keras/adversarial_regularization.py

@@ -20,6 +20,7 @@
 import collections
 import functools
 import types
+from keras.utils import traceback_utils
 
 import attr
 import neural_structured_learning.configs as nsl_configs
@@ -699,6 +700,61 @@ def save(self, *args, **kwargs):
         'Saving `AdversarialRegularization` models is currently not supported. '
         'Consider using `save_weights` or saving the `base_model`.')
 
+
+  @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, filepath, overwrite, include_optimizer, save_format,
+                    signatures, options, save_traces)
+
+
+
+
+
+
   def perturb_on_batch(self, x, **config_kwargs):
     """Perturbs the given input to generates adversarial examples.