Only pass defined parameters to base Estimator's model_fn.

PiperOrigin-RevId: 315571871

Author: csferng

neural_structured_learning/estimator/adversarial_regularization.py

@@ -18,9 +18,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+import inspect
+
 import neural_structured_learning.configs as nsl_configs
 import neural_structured_learning.lib as nsl_lib
-
 import tensorflow as tf
 
 
@@ -55,6 +57,10 @@ def add_adversarial_regularization(estimator,
     adv_config = nsl_configs.AdvRegConfig()
 
   base_model_fn = estimator._model_fn  # pylint: disable=protected-access
+  try:
+    base_model_fn_args = inspect.signature(base_model_fn).parameters.keys()
+  except AttributeError:  # For Python 2 compatibility
+    base_model_fn_args = inspect.getargspec(base_model_fn).args  # pylint: disable=deprecated-method
 
   def adv_model_fn(features, labels, mode, params=None, config=None):
     """The adversarial-regularized model_fn.
@@ -82,19 +88,22 @@ def adv_model_fn(features, labels, mode, params=None, config=None):
     Returns:
       A `tf.estimator.EstimatorSpec` with adversarial regularization.
     """
+    # Parameters 'params' and 'config' are optional. If they are not passed,
+    # then it is possible for base_model_fn not to accept these arguments.
+    # See documentation for tf.estimator.Estimator for additional context.
+    kwargs = {'mode': mode}
+    if 'params' in base_model_fn_args:
+      kwargs['params'] = params
+    if 'config' in base_model_fn_args:
+      kwargs['config'] = config
+    base_fn = functools.partial(base_model_fn, **kwargs)
 
     # Uses the same variable scope for calculating the original objective and
     # adversarial regularization.
     with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(),
                                      reuse=tf.compat.v1.AUTO_REUSE,
                                      auxiliary_name_scope=False):
-      # If no 'params' is passed, then it is possible for base_model_fn not to
-      # accept a 'params' argument. See documentation for tf.estimator.Estimator
-      # for additional context.
-      base_args = [mode, params, config] if params else [mode, config]
-      spec_fn = lambda feature, label: base_model_fn(feature, label, *base_args)
-
-      original_spec = spec_fn(features, labels)
+      original_spec = base_fn(features, labels)
 
       # Adversarial regularization only happens in training.
       if mode != tf.estimator.ModeKeys.TRAIN:
@@ -107,11 +116,11 @@ def adv_model_fn(features, labels, mode, params=None, config=None):
           # The pgd_model_fn is a dummy identity function since loss is
           # directly available from spec_fn.
           pgd_model_fn=lambda features: features,
-          pgd_loss_fn=lambda labels, features: spec_fn(features, labels).loss,
+          pgd_loss_fn=lambda labels, features: base_fn(features, labels).loss,
           pgd_labels=labels)
 
       # Runs the base model again to compute loss on adv_neighbor.
-      adv_spec = spec_fn(adv_neighbor, labels)
+      adv_spec = base_fn(adv_neighbor, labels)
 
       final_loss = original_spec.loss + adv_config.multiplier * adv_spec.loss
 

neural_structured_learning/estimator/graph_regularization.py

@@ -17,6 +17,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
+
 import neural_structured_learning.configs as configs
 from neural_structured_learning.lib import distances
 from neural_structured_learning.lib import utils
@@ -51,6 +53,10 @@ def add_graph_regularization(estimator,
     graph_reg_config = configs.GraphRegConfig()
 
   base_model_fn = estimator._model_fn  # pylint: disable=protected-access
+  try:
+    base_model_fn_args = inspect.signature(base_model_fn).parameters.keys()
+  except AttributeError:  # For Python 2 compatibility
+    base_model_fn_args = inspect.getargspec(base_model_fn).args  # pylint: disable=deprecated-method
 
   def graph_reg_model_fn(features, labels, mode, params=None, config=None):
     """The graph-regularized model function.
@@ -79,6 +85,14 @@ def graph_reg_model_fn(features, labels, mode, params=None, config=None):
     Returns:
       A `tf.estimator.EstimatorSpec` with graph regularization.
     """
+    # Parameters 'params' and 'config' are optional. If they are not passed,
+    # then it is possible for base_model_fn not to accept these arguments.
+    # See documentation for tf.estimator.Estimator for additional context.
+    kwargs = {'mode': mode}
+    if 'params' in base_model_fn_args:
+      kwargs['params'] = params
+    if 'config' in base_model_fn_args:
+      kwargs['config'] = config
 
     # Uses the same variable scope for calculating the original objective and
     # the graph regularization loss term.
@@ -100,13 +114,7 @@ def graph_reg_model_fn(features, labels, mode, params=None, config=None):
         sample_features = utils.strip_neighbor_features(
             features, graph_reg_config.neighbor_config)
 
-      # If no 'params' is passed, then it is possible for base_model_fn not to
-      # accept a 'params' argument. See documentation for tf.estimator.Estimator
-      # for additional context.
-      if params:
-        base_spec = base_model_fn(sample_features, labels, mode, params, config)
-      else:
-        base_spec = base_model_fn(sample_features, labels, mode, config)
+      base_spec = base_model_fn(sample_features, labels, **kwargs)
 
       has_nbr_inputs = nbr_weights is not None and nbr_features