Fix batch stat updates for adversarial-regularized estimator in TensorFlow 1.x

PiperOrigin-RevId: 318535128

Author: csferng

neural_structured_learning/estimator/adversarial_regularization.py

@@ -130,10 +130,15 @@ def adv_model_fn(features, labels, mode, params=None, config=None):
       else:
         optimizer = optimizer_fn()
 
-      final_train_op = optimizer.minimize(
+      train_op = optimizer.minimize(
           loss=final_loss, global_step=tf.compat.v1.train.get_global_step())
 
-    return original_spec._replace(loss=final_loss, train_op=final_train_op)
+      update_ops = tf.compat.v1.get_collection(
+          tf.compat.v1.GraphKeys.UPDATE_OPS)
+      if update_ops:
+        train_op = tf.group(train_op, *update_ops)
+
+    return original_spec._replace(loss=final_loss, train_op=train_op)
 
   # Replaces the model_fn while keeps other fields/methods in the estimator.
   estimator._model_fn = adv_model_fn  # pylint: disable=protected-access

neural_structured_learning/estimator/adversarial_regularization_test.py

@@ -126,6 +126,27 @@ def test_adversarial_wrapper_adds_regularization(self, adv_step_size,
     self.assertAllClose(new_bias, adv_est.get_variable_value(BIAS_VARIABLE))
     self.assertAllClose(new_weight, adv_est.get_variable_value(WEIGHT_VARIABLE))
 
+  @test_util.run_v1_only('Requires tf.train.GradientDescentOptimizer')
+  def test_adversarial_wrapper_saving_batch_statistics(self):
+    x0, y0 = np.array([[0.9, 0.1], [0.2, 0.8]]), np.array([1, 0])
+    input_fn = single_batch_input_fn({FEATURE_NAME: x0}, y0)
+    fc = tf.feature_column.numeric_column(FEATURE_NAME, shape=[2])
+    base_est = tf.estimator.DNNClassifier(
+        hidden_units=[4],
+        feature_columns=[fc],
+        model_dir=self.model_dir,
+        batch_norm=True)
+    adv_est = nsl_estimator.add_adversarial_regularization(
+        base_est,
+        optimizer_fn=lambda: tf.train.GradientDescentOptimizer(0.005))
+
+    adv_est.train(input_fn=input_fn, steps=1)
+    moving_mean = adv_est.get_variable_value(
+        'dnn/hiddenlayer_0/batchnorm_0/moving_mean')
+    moving_variance = adv_est.get_variable_value(
+        'dnn/hiddenlayer_0/batchnorm_0/moving_variance')
+    self.assertNotAllClose(moving_mean, np.zeros(moving_mean.shape))
+    self.assertNotAllClose(moving_variance, np.ones(moving_variance.shape))
 
 if __name__ == '__main__':
   tf.test.main()