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: 441225284

Author: hertschuh

classifier_utils.py

@@ -25,6 +25,7 @@
 from albert import optimization
 from albert import tokenization
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import metrics as contrib_metrics
 from tensorflow.contrib import tpu as contrib_tpu
@@ -835,7 +836,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     else:
       is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, probabilities, logits, predictions) = \
         create_model(albert_config, is_training, input_ids, input_mask,
@@ -867,7 +868,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
 
       train_op = optimization.create_optimizer(
           total_loss, learning_rate, num_train_steps, num_warmup_steps,
@@ -878,7 +879,7 @@ def tpu_scaffold():
           loss=total_loss,
           train_op=train_op,
           scaffold_fn=scaffold_fn)
-    elif mode == tf.estimator.ModeKeys.EVAL:
+    elif mode == tf_estimator.ModeKeys.EVAL:
       if task_name not in ["sts-b", "cola"]:
         def metric_fn(per_example_loss, label_ids, logits, is_real_example):
           predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)

race_utils.py

@@ -27,6 +27,7 @@
 from albert import optimization
 from albert import tokenization
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import tpu as contrib_tpu
 
 
@@ -356,7 +357,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     else:
       is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, probabilities, logits, predictions) = \
         create_model(albert_config, is_training, input_ids, input_mask,
@@ -389,7 +390,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
 
       train_op = optimization.create_optimizer(
           total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)
@@ -399,7 +400,7 @@ def tpu_scaffold():
           loss=total_loss,
           train_op=train_op,
           scaffold_fn=scaffold_fn)
-    elif mode == tf.estimator.ModeKeys.EVAL:
+    elif mode == tf_estimator.ModeKeys.EVAL:
       def metric_fn(per_example_loss, label_ids, logits, is_real_example):
         predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
         accuracy = tf.metrics.accuracy(

run_classifier.py

@@ -25,6 +25,7 @@
 from albert import fine_tuning_utils
 from albert import modeling
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import cluster_resolver as contrib_cluster_resolver
 from tensorflow.contrib import tpu as contrib_tpu
 
@@ -177,7 +178,7 @@ def _serving_input_receiver_fn():
       t = tf.to_int32(t)
     feature_map[name] = t
 
-  return tf.estimator.export.ServingInputReceiver(
+  return tf_estimator.export.ServingInputReceiver(
       features=feature_map, receiver_tensors=serialized_example)
 
 

run_pretraining.py

@@ -24,6 +24,7 @@
 from albert import optimization
 from six.moves import range
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import cluster_resolver as contrib_cluster_resolver
 from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import tpu as contrib_tpu
@@ -153,7 +154,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     # it does represent sentence_order_labels.
     sentence_order_labels = features["next_sentence_labels"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     model = modeling.AlbertModel(
         config=albert_config,
@@ -217,7 +218,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       train_op = optimization.create_optimizer(
           total_loss, learning_rate, num_train_steps, num_warmup_steps,
           use_tpu, optimizer, poly_power, start_warmup_step)
@@ -227,7 +228,7 @@ def tpu_scaffold():
           loss=total_loss,
           train_op=train_op,
           scaffold_fn=scaffold_fn)
-    elif mode == tf.estimator.ModeKeys.EVAL:
+    elif mode == tf_estimator.ModeKeys.EVAL:
 
       def metric_fn(*args):
         """Computes the loss and accuracy of the model."""

run_squad_v1.py

@@ -29,6 +29,7 @@
 from albert import squad_utils
 import six
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.contrib import cluster_resolver as contrib_cluster_resolver
 from tensorflow.contrib import tpu as contrib_tpu
@@ -236,7 +237,7 @@ def _seq_serving_input_fn():
         "input_mask": input_mask,
         "segment_ids": segment_ids
     }
-    return tf.estimator.export.ServingInputReceiver(features=inputs,
+    return tf_estimator.export.ServingInputReceiver(features=inputs,
                                                     receiver_tensors=inputs)
 
   return _seq_serving_input_fn

squad_utils.py

@@ -33,6 +33,7 @@
 from six.moves import map
 from six.moves import range
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import layers as contrib_layers
 from tensorflow.contrib import tpu as contrib_tpu
@@ -767,7 +768,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     input_mask = features["input_mask"]
     segment_ids = features["segment_ids"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (start_logits, end_logits) = create_v1_model(
         albert_config=albert_config,
@@ -809,7 +810,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       seq_length = modeling.get_shape_list(input_ids)[1]
 
       def compute_loss(logits, positions):
@@ -836,7 +837,7 @@ def compute_loss(logits, positions):
           loss=total_loss,
           train_op=train_op,
           scaffold_fn=scaffold_fn)
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {
           "start_log_prob": start_logits,
           "end_log_prob": end_logits,
@@ -1594,7 +1595,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     input_mask = features["input_mask"]
     segment_ids = features["segment_ids"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     outputs = create_v2_model(
         albert_config=albert_config,
@@ -1636,7 +1637,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       seq_length = modeling.get_shape_list(input_ids)[1]
 
       def compute_loss(log_probs, positions):
@@ -1671,7 +1672,7 @@ def compute_loss(log_probs, positions):
           loss=total_loss,
           train_op=train_op,
           scaffold_fn=scaffold_fn)
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {
           "unique_ids": features["unique_ids"],
           "start_top_index": outputs["start_top_index"],