Make the adversarial estimator test more robust.

PiperOrigin-RevId: 326539266

Author: csferng

neural_structured_learning/estimator/adversarial_regularization_test.py

@@ -128,13 +128,16 @@ def test_adversarial_wrapper_adds_regularization(self, adv_step_size,
 
   @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])
+    x0 = np.array([[0.9, 0.1], [0.2, -0.8], [-0.7, -0.3], [-0.4, 0.6]])
+    y0 = np.array([1, 0, 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,
+        activation_fn=lambda x: tf.abs(x) + 0.1,
+        dropout=None,
         batch_norm=True)
     adv_est = nsl_estimator.add_adversarial_regularization(
         base_est,
@@ -145,6 +148,8 @@ def test_adversarial_wrapper_saving_batch_statistics(self):
         'dnn/hiddenlayer_0/batchnorm_0/moving_mean')
     moving_variance = adv_est.get_variable_value(
         'dnn/hiddenlayer_0/batchnorm_0/moving_variance')
+    # The activation function always returns a positive number, so the batch
+    # mean cannot be zero if updated successfully.
     self.assertNotAllClose(moving_mean, np.zeros(moving_mean.shape))
     self.assertNotAllClose(moving_variance, np.ones(moving_variance.shape))