Explicitly import estimator from tensorflow as a separate import instead of accessing it via tf.estimator and depend on the tensorflow estimator target.

PiperOrigin-RevId: 436596645

Author: qlzh727

neural_structured_learning/estimator/BUILD

@@ -40,6 +40,7 @@ py_library(
         "//neural_structured_learning/configs",
         "//neural_structured_learning/lib",
         # package tensorflow
+        # package tensorflow:tensorflow_estimator
     ],
 )
 
@@ -54,6 +55,7 @@ py_test(
         "//neural_structured_learning/configs",
         # package numpy
         # package tensorflow
+        # package tensorflow:tensorflow_estimator
         # package tensorflow framework_test_lib,
     ],
 )
@@ -67,6 +69,7 @@ py_library(
         "//neural_structured_learning/lib:distances",
         "//neural_structured_learning/lib:utils",
         # package tensorflow
+        # package tensorflow:tensorflow_estimator
     ],
 )
 
@@ -81,6 +84,7 @@ py_test(
         "//neural_structured_learning/configs",
         # package numpy
         # package tensorflow
+        # package tensorflow:tensorflow_estimator
         # package tensorflow framework_test_lib,
     ],
 )

neural_structured_learning/estimator/adversarial_regularization.py

@@ -24,6 +24,7 @@
 import neural_structured_learning.configs as nsl_configs
 import neural_structured_learning.lib as nsl_lib
 import tensorflow as tf
+from tensorflow import estimator as tf_estimator
 
 
 def add_adversarial_regularization(estimator,
@@ -106,7 +107,7 @@ def adv_model_fn(features, labels, mode, params=None, config=None):
       original_spec = base_fn(features, labels)
 
       # Adversarial regularization only happens in training.
-      if mode != tf.estimator.ModeKeys.TRAIN:
+      if mode != tf_estimator.ModeKeys.TRAIN:
         return original_spec
 
       adv_neighbor, _ = nsl_lib.gen_adv_neighbor(

neural_structured_learning/estimator/adversarial_regularization_test.py

@@ -26,6 +26,7 @@
 import neural_structured_learning.estimator as nsl_estimator
 import numpy as np
 import tensorflow as tf
+from tensorflow import estimator as tf_estimator
 
 from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
 
@@ -66,7 +67,7 @@ def build_linear_regressor(self, weight, bias):
 
     fc = tf.feature_column.numeric_column(FEATURE_NAME,
                                           shape=np.array(weight).shape)
-    return tf.estimator.LinearRegressor(
+    return tf_estimator.LinearRegressor(
         feature_columns=(fc,), model_dir=self.model_dir, optimizer='SGD')
 
   @test_util.run_v1_only('Requires tf.GraphKeys')
@@ -132,7 +133,7 @@ def test_adversarial_wrapper_saving_batch_statistics(self):
     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(
+    base_est = tf_estimator.DNNClassifier(
         hidden_units=[4],
         feature_columns=[fc],
         model_dir=self.model_dir,

neural_structured_learning/estimator/graph_regularization.py

@@ -24,6 +24,7 @@
 from neural_structured_learning.lib import utils
 
 import tensorflow as tf
+from tensorflow import estimator as tf_estimator
 
 
 def add_graph_regularization(estimator,
@@ -108,7 +109,7 @@ def graph_reg_model_fn(features, labels, mode, params=None, config=None):
         auxiliary_name_scope=False):
       nbr_features = dict()
       nbr_weights = None
-      if mode == tf.estimator.ModeKeys.TRAIN:
+      if mode == tf_estimator.ModeKeys.TRAIN:
         # Extract sample features, neighbor features, and neighbor weights if we
         # are in training mode.
         sample_features, nbr_features, nbr_weights = (
@@ -130,7 +131,7 @@ def graph_reg_model_fn(features, labels, mode, params=None, config=None):
       # - neighbor inputs exist
       # - the graph regularization multiplier is greater than zero.
       # So, return early if any of these conditions is false.
-      if (not has_nbr_inputs or mode != tf.estimator.ModeKeys.TRAIN or
+      if (not has_nbr_inputs or mode != tf_estimator.ModeKeys.TRAIN or
           graph_reg_config.multiplier <= 0):
         return base_spec
 

neural_structured_learning/estimator/graph_regularization_test.py

@@ -26,6 +26,7 @@
 
 import numpy as np
 import tensorflow as tf
+from tensorflow import estimator as tf_estimator
 
 from google.protobuf import text_format
 from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
@@ -136,7 +137,7 @@ def build_linear_regressor(self, weight, weight_shape, bias, bias_shape):
 
     fc = tf.feature_column.numeric_column(
         FEATURE_NAME, shape=np.array(weight).shape)
-    return tf.estimator.LinearRegressor(
+    return tf_estimator.LinearRegressor(
         feature_columns=(fc,), model_dir=self.model_dir, optimizer='SGD')
 
   @test_util.run_v1_only('Requires tf.get_variable')
@@ -452,7 +453,7 @@ def embedding_fn(features, mode, params=None):
         # The always-positive activation funciton is to make sure the batch mean
         # is non-zero.
         batch_norm_layer = tf.compat.v1.layers.batch_normalization(
-            hidden_layer, training=(mode == tf.estimator.ModeKeys.TRAIN))
+            hidden_layer, training=(mode == tf_estimator.ModeKeys.TRAIN))
       return batch_norm_layer
 
     def model_fn(features, labels, mode, params=None, config=None):
@@ -461,25 +462,25 @@ def model_fn(features, labels, mode, params=None, config=None):
       with tf.compat.v1.variable_scope('logit', reuse=tf.AUTO_REUSE):
         logits = tf.compat.v1.layers.dense(embeddings, units=1)
       predictions = tf.argmax(logits, 1)
-      if mode == tf.estimator.ModeKeys.PREDICT:
-        return tf.estimator.EstimatorSpec(
+      if mode == tf_estimator.ModeKeys.PREDICT:
+        return tf_estimator.EstimatorSpec(
             mode=mode,
             predictions={
                 'logits': logits,
                 'predictions': predictions
             })
 
       loss = tf.losses.sigmoid_cross_entropy(labels, logits)
-      if mode == tf.estimator.ModeKeys.EVAL:
-        return tf.estimator.EstimatorSpec(mode=mode, loss=loss)
+      if mode == tf_estimator.ModeKeys.EVAL:
+        return tf_estimator.EstimatorSpec(mode=mode, loss=loss)
 
       optimizer = optimizer_fn()
       train_op = optimizer.minimize(
           loss, global_step=tf.compat.v1.train.get_global_step())
       update_ops = tf.compat.v1.get_collection(
           tf.compat.v1.GraphKeys.UPDATE_OPS)
       train_op = tf.group(train_op, *update_ops)
-      return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
+      return tf_estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
 
     def input_fn():
       nbr_feature = '{}{}_{}'.format(NBR_FEATURE_PREFIX, 0, FEATURE_NAME)
@@ -492,7 +493,7 @@ def input_fn():
       labels = tf.constant([[1], [0], [1]])
       return tf.data.Dataset.from_tensor_slices((features, labels)).batch(3)
 
-    base_est = tf.estimator.Estimator(
+    base_est = tf_estimator.Estimator(
         model_fn, model_dir=self.model_dir, params=None)
     graph_reg_config = nsl_configs.make_graph_reg_config(
         max_neighbors=1, multiplier=1)