Adds model_kwargs argument to nsl.keras.adversarial_loss.

The argument is for passing additional arguments to the base model, so
that some layers depending on such argument can function properly.
For example, `BatchNormalization` layer requires `training=True` during
training phase.

Fixes #28

PiperOrigin-RevId: 273618791

Author: csferng

neural_structured_learning/keras/adversarial_regularization.py

@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import collections
+import functools
 import types
 
 import attr
@@ -36,7 +37,8 @@ def adversarial_loss(features,
                      adv_config=None,
                      predictions=None,
                      labeled_loss=None,
-                     gradient_tape=None):
+                     gradient_tape=None,
+                     model_kwargs=None):
   """Computes the adversarial loss for `model` given `features` and `labels`.
 
   This utility function adds adversarial perturbations to the input `features`,
@@ -109,6 +111,8 @@ def adversarial_loss(features,
       adversarial regularization. In eager mode, the `gradient_tape` has to be
       set as well.
     gradient_tape: (optional) A `tf.GradientTape` object watching `features`.
+    model_kwargs: (optional) A dictionary of additional keyword arguments to be
+      passed to the `model`.
 
   Returns:
     A `Tensor` for adversarial regularization loss, i.e. labeled loss on
@@ -118,6 +122,9 @@ def adversarial_loss(features,
   if adv_config is None:
     adv_config = nsl_configs.AdvRegConfig()
 
+  if model_kwargs is not None:
+    model = functools.partial(model, **model_kwargs)
+
   # Calculates labeled_loss if not provided.
   if labeled_loss is None:
     # Reuses the tape if provided; otherwise creates a new tape.
@@ -626,16 +633,16 @@ def call(self, inputs, **kwargs):
       self.add_metric(value, aggregation=aggregation, name=name)
 
     # Adversarial loss.
-    base_model_fn = lambda inputs: self.base_model(inputs, **kwargs)
     adv_loss = adversarial_loss(
         inputs,
         labels,
-        base_model_fn,
+        self.base_model,
         self._compute_total_loss,
         sample_weights=sample_weights,
         adv_config=self.adv_config,
         labeled_loss=labeled_loss,
-        gradient_tape=tape)
+        gradient_tape=tape,
+        model_kwargs=kwargs)
     self.add_loss(self.adv_config.multiplier * adv_loss)
     self.add_metric(adv_loss, name='adversarial_loss', aggregation='mean')
     return outputs

neural_structured_learning/keras/adversarial_regularization_test.py

@@ -202,6 +202,23 @@ def wrapped_loss_fn(*args, **kwargs):
     self.assertEqual(1, call_count['model'])
     self.assertEqual(1, call_count['loss_fn'])
 
+  def test_with_model_kwargs(self):
+    w = np.array([[4.0], [-3.0]])
+    x0 = np.array([[2.0, 3.0]])
+    y0 = np.array([[0.0]])
+    model = build_linear_keras_sequential_model(input_shape=(2,), weights=w)
+    model.add(tf.keras.layers.BatchNormalization())
+
+    adv_loss = adversarial_regularization.adversarial_loss(
+        features={'feature': tf.constant(x0)},
+        labels=tf.constant(y0),
+        model=model,
+        loss_fn=keras.losses.MeanSquaredError(),
+        adv_config=self.adv_config,
+        model_kwargs={'training': True})
+    # BatchNormalization returns 0 for signle-example batch when training=True.
+    self.assertAllClose(0.0, self.evaluate(adv_loss))
+
 
 class AdversarialRegularizationTest(tf.test.TestCase, parameterized.TestCase):