Explicitly import estimator from tensorflow as a separate import inst…

…ead of

accessing it via tf.estimator and depend on the tensorflow estimator target.

PiperOrigin-RevId: 437339196

language/bert_extraction/steal_bert_classifier/models/run_classifier.py

@@ -22,6 +22,7 @@
 from bert import optimization
 from bert import tokenization
 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
@@ -834,7 +835,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, logits,
      probabilities) = create_model(bert_config, is_training, input_ids,
@@ -866,7 +867,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,
@@ -877,7 +878,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)

language/bert_extraction/steal_bert_classifier/models/run_classifier_distillation.py

@@ -22,6 +22,7 @@
 from bert import optimization
 from bert import tokenization
 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
@@ -580,7 +581,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, logits,
      probabilities) = create_model(bert_config, is_training, input_ids,
@@ -612,7 +613,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,
@@ -623,7 +624,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)

language/bert_extraction/steal_bert_classifier/models/run_classifier_membership.py

@@ -21,6 +21,7 @@
 from bert import modeling
 from bert import tokenization
 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
@@ -611,7 +612,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)
 
     membership_logits, membership_vars = create_model(bert_config, is_training,
                                                       input_ids, input_mask,
@@ -647,7 +648,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
 
       one_hot_positions = tf.one_hot(label_ids, depth=2, dtype=tf.float32)
       per_example_loss = -tf.reduce_sum(
@@ -670,7 +671,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:
 
       one_hot_positions = tf.one_hot(label_ids, depth=2, dtype=tf.float32)
       per_example_loss = -tf.reduce_sum(

language/bert_extraction/steal_bert_qa/models/run_bert_boolq.py

@@ -27,6 +27,7 @@
 from bert import tokenization
 import numpy as np
 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
@@ -454,7 +455,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     segment_ids = features["segment_ids"]
     label_ids = features["label_ids"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, logits, probabilities) = create_model(
         bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
@@ -487,7 +488,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)
@@ -497,7 +498,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):
         predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
@@ -515,7 +516,7 @@ def metric_fn(per_example_loss, label_ids, logits):
           eval_metrics=eval_metrics,
           scaffold_fn=scaffold_fn)
 
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {"probabilities": probabilities}
       output_spec = contrib_tpu.TPUEstimatorSpec(
           mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)

language/bert_extraction/steal_bert_qa/models/run_bert_boolq_distill.py

@@ -27,6 +27,7 @@
 from bert import tokenization
 import numpy as np
 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
@@ -471,7 +472,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     label_ids = features["label_ids"]
     gt_probs = features["probs"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, logits,
      probabilities) = create_model(bert_config, is_training, input_ids,
@@ -505,7 +506,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)
@@ -515,7 +516,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):
         predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
@@ -533,7 +534,7 @@ def metric_fn(per_example_loss, label_ids, logits):
           eval_metrics=eval_metrics,
           scaffold_fn=scaffold_fn)
 
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {"probabilities": probabilities}
       output_spec = contrib_tpu.TPUEstimatorSpec(
           mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)

language/bert_extraction/steal_bert_qa/models/run_squad.py

@@ -29,6 +29,7 @@
 import six
 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
@@ -616,7 +617,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_model(
         bert_config=bert_config,
@@ -652,7 +653,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):
@@ -680,7 +681,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 = {
           "unique_ids": unique_ids,
           "start_logits": start_logits,

language/bert_extraction/steal_bert_qa/models/run_squad_membership.py

@@ -27,6 +27,7 @@
 import six
 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
@@ -575,7 +576,7 @@ def tpu_scaffold():
                       init_string)
 
     output_spec = None
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
 
       one_hot_positions = tf.one_hot(label_ids, depth=2, dtype=tf.float32)
       loss = -tf.reduce_mean(
@@ -593,7 +594,7 @@ def tpu_scaffold():
       output_spec = contrib_tpu.TPUEstimatorSpec(
           mode=mode, loss=loss, train_op=train_op, scaffold_fn=scaffold_fn)
 
-    elif mode == tf.estimator.ModeKeys.EVAL:
+    elif mode == tf_estimator.ModeKeys.EVAL:
 
       one_hot_positions = tf.one_hot(label_ids, depth=2, dtype=tf.float32)
       per_example_loss = -1 * tf.reduce_sum(
@@ -617,7 +618,7 @@ def metric_fn(per_example_loss, label_ids, membership_logits):
           eval_metrics=eval_metrics,
           scaffold_fn=scaffold_fn)
 
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {
           "unique_ids": unique_ids,
           "membership_probs": membership_probs

language/bert_extraction/steal_bert_qa/utils/run_bert_boolq_diff.py

@@ -28,6 +28,7 @@
 import numpy as np
 from six.moves import zip
 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
@@ -434,7 +435,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     label_ids = features["label_ids"]
     gt_probs = features["probs"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, logits,
      probabilities) = create_model(bert_config, is_training, input_ids,
@@ -468,7 +469,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)
@@ -478,7 +479,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):
         predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
@@ -496,7 +497,7 @@ def metric_fn(per_example_loss, label_ids, logits):
           eval_metrics=eval_metrics,
           scaffold_fn=scaffold_fn)
 
-    elif mode == tf.estimator.ModeKeys.PREDICT:
+    elif mode == tf_estimator.ModeKeys.PREDICT:
       predictions = {"probabilities": probabilities}
       output_spec = contrib_tpu.TPUEstimatorSpec(
           mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)

language/boolq/run_bert_boolq.py

@@ -27,6 +27,7 @@
 from bert import tokenization
 import numpy as np
 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
 
@@ -438,7 +439,7 @@ def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
     segment_ids = features["segment_ids"]
     label_ids = features["label_ids"]
 
-    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
+    is_training = (mode == tf_estimator.ModeKeys.TRAIN)
 
     (total_loss, per_example_loss, logits, probabilities) = create_model(
         bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
@@ -471,7 +472,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)
@@ -481,7 +482,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):
         predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)

language/boolq/run_recurrent_model_boolq.py

@@ -31,6 +31,7 @@
 from language.common.utils import tensor_utils
 import numpy as np
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 from tensorflow.contrib import lookup as contrib_lookup
 
 # Dataset parameters, these need to be pointed at the appropriate targets
@@ -214,7 +215,7 @@ def apply_lstm(x, seq_len):
       hidden_size=FLAGS.lstm_dim,
       num_layers=1,
       dropout_ratio=0.0,
-      mode=tf.estimator.ModeKeys.TRAIN,
+      mode=tf_estimator.ModeKeys.TRAIN,
       use_cudnn=None)
 
 
@@ -432,14 +433,14 @@ def train():
   # with named arguments, so pylint: disable=unused-argument
   def model_function(features, labels, mode, params):
     """Builds the `tf.estimator.EstimatorSpec` to train/eval with."""
-    is_train = mode == tf.estimator.ModeKeys.TRAIN
+    is_train = mode == tf_estimator.ModeKeys.TRAIN
     logits = predict(is_train, embeddings, features["premise"],
                      features["hypothesis"])
 
     loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
         labels=tf.to_int32(labels), logits=logits)
     loss = tf.reduce_mean(loss)
-    if mode == tf.estimator.ModeKeys.TRAIN:
+    if mode == tf_estimator.ModeKeys.TRAIN:
       train_op = get_train_op(loss)
     else:
       # Don't build the train_op unnecessarily, since the ADAM variables can
@@ -460,7 +461,7 @@ def _init_fn(_, sess):
 
       scaffold = tf.train.Scaffold(init_fn=_init_fn)
 
-    return tf.estimator.EstimatorSpec(
+    return tf_estimator.EstimatorSpec(
         mode=mode,
         scaffold=scaffold,
         loss=loss,

language/boolq/utils/best_checkpoint_exporter.py

@@ -20,9 +20,10 @@
 
 import os
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 
-class BestCheckpointExporter(tf.estimator.Exporter):
+class BestCheckpointExporter(tf_estimator.Exporter):
   """Exporter that saves the model's best checkpoint.
 
   We use this over `tf.estimator.BestExporter` since we don't want to