From 721f26e1b7301ee36976c9705bc4effecd0fe9e5 Mon Sep 17 00:00:00 2001 From: Fabien Hertschuh Date: Fri, 8 Apr 2022 12:32:13 -0700 Subject: [PATCH] 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: 440427683 --- classifier_utils.py | 7 ++++--- race_utils.py | 7 ++++--- run_classifier.py | 3 ++- run_pretraining.py | 7 ++++--- run_squad_v1.py | 3 ++- squad_utils.py | 13 +++++++------ 6 files changed, 23 insertions(+), 17 deletions(-) diff --git a/classifier_utils.py b/classifier_utils.py index 2d997826..0edcbd70 100644 --- a/classifier_utils.py +++ b/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) diff --git a/race_utils.py b/race_utils.py index 4c010c12..a88e8785 100644 --- a/race_utils.py +++ b/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( diff --git a/run_classifier.py b/run_classifier.py index 2d3a1e03..d1b7a5d4 100644 --- a/run_classifier.py +++ b/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) diff --git a/run_pretraining.py b/run_pretraining.py index 52f2fa9b..6c8bbc64 100644 --- a/run_pretraining.py +++ b/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.""" diff --git a/run_squad_v1.py b/run_squad_v1.py index 8cfbb6b8..132e9bf0 100644 --- a/run_squad_v1.py +++ b/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 diff --git a/squad_utils.py b/squad_utils.py index 458c507c..a5f6cb28 100644 --- a/squad_utils.py +++ b/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"],