23
23
import os
24
24
25
25
import tensorflow .compat .v1 as tf
26
+ from tensorflow .compat .v1 import estimator as tf_estimator
26
27
import tensorflow .compat .v2 as tf2
27
28
import tf_slim as slim
28
29
@@ -465,7 +466,7 @@ def model_fn(features, labels, mode, params=None):
465
466
"""
466
467
params = params or {}
467
468
total_loss , train_op , detections , export_outputs = None , None , None , None
468
- is_training = mode == tf . estimator .ModeKeys .TRAIN
469
+ is_training = mode == tf_estimator .ModeKeys .TRAIN
469
470
470
471
# Make sure to set the Keras learning phase. True during training,
471
472
# False for inference.
@@ -479,11 +480,11 @@ def model_fn(features, labels, mode, params=None):
479
480
is_training = is_training , add_summaries = (not use_tpu ))
480
481
scaffold_fn = None
481
482
482
- if mode == tf . estimator .ModeKeys .TRAIN :
483
+ if mode == tf_estimator .ModeKeys .TRAIN :
483
484
labels = unstack_batch (
484
485
labels ,
485
486
unpad_groundtruth_tensors = train_config .unpad_groundtruth_tensors )
486
- elif mode == tf . estimator .ModeKeys .EVAL :
487
+ elif mode == tf_estimator .ModeKeys .EVAL :
487
488
# For evaling on train data, it is necessary to check whether groundtruth
488
489
# must be unpadded.
489
490
boxes_shape = (
@@ -493,7 +494,7 @@ def model_fn(features, labels, mode, params=None):
493
494
labels = unstack_batch (
494
495
labels , unpad_groundtruth_tensors = unpad_groundtruth_tensors )
495
496
496
- if mode in (tf . estimator . ModeKeys .TRAIN , tf . estimator .ModeKeys .EVAL ):
497
+ if mode in (tf_estimator . ModeKeys .TRAIN , tf_estimator .ModeKeys .EVAL ):
497
498
provide_groundtruth (detection_model , labels )
498
499
499
500
preprocessed_images = features [fields .InputDataFields .image ]
@@ -514,7 +515,7 @@ def model_fn(features, labels, mode, params=None):
514
515
def postprocess_wrapper (args ):
515
516
return detection_model .postprocess (args [0 ], args [1 ])
516
517
517
- if mode in (tf . estimator . ModeKeys .EVAL , tf . estimator .ModeKeys .PREDICT ):
518
+ if mode in (tf_estimator . ModeKeys .EVAL , tf_estimator .ModeKeys .PREDICT ):
518
519
if use_tpu and postprocess_on_cpu :
519
520
detections = tf .tpu .outside_compilation (
520
521
postprocess_wrapper ,
@@ -525,7 +526,7 @@ def postprocess_wrapper(args):
525
526
prediction_dict ,
526
527
features [fields .InputDataFields .true_image_shape ]))
527
528
528
- if mode == tf . estimator .ModeKeys .TRAIN :
529
+ if mode == tf_estimator .ModeKeys .TRAIN :
529
530
load_pretrained = hparams .load_pretrained if hparams else False
530
531
if train_config .fine_tune_checkpoint and load_pretrained :
531
532
if not train_config .fine_tune_checkpoint_type :
@@ -557,8 +558,8 @@ def tpu_scaffold():
557
558
tf .train .init_from_checkpoint (train_config .fine_tune_checkpoint ,
558
559
available_var_map )
559
560
560
- if mode in (tf . estimator . ModeKeys .TRAIN , tf . estimator .ModeKeys .EVAL ):
561
- if (mode == tf . estimator .ModeKeys .EVAL and
561
+ if mode in (tf_estimator . ModeKeys .TRAIN , tf_estimator .ModeKeys .EVAL ):
562
+ if (mode == tf_estimator .ModeKeys .EVAL and
562
563
eval_config .use_dummy_loss_in_eval ):
563
564
total_loss = tf .constant (1.0 )
564
565
losses_dict = {'Loss/total_loss' : total_loss }
@@ -590,7 +591,7 @@ def tpu_scaffold():
590
591
training_optimizer , optimizer_summary_vars = optimizer_builder .build (
591
592
train_config .optimizer )
592
593
593
- if mode == tf . estimator .ModeKeys .TRAIN :
594
+ if mode == tf_estimator .ModeKeys .TRAIN :
594
595
if use_tpu :
595
596
training_optimizer = tf .tpu .CrossShardOptimizer (training_optimizer )
596
597
@@ -628,16 +629,16 @@ def tpu_scaffold():
628
629
summaries = summaries ,
629
630
name = '' ) # Preventing scope prefix on all variables.
630
631
631
- if mode == tf . estimator .ModeKeys .PREDICT :
632
+ if mode == tf_estimator .ModeKeys .PREDICT :
632
633
exported_output = exporter_lib .add_output_tensor_nodes (detections )
633
634
export_outputs = {
634
635
tf .saved_model .signature_constants .PREDICT_METHOD_NAME :
635
- tf . estimator .export .PredictOutput (exported_output )
636
+ tf_estimator .export .PredictOutput (exported_output )
636
637
}
637
638
638
639
eval_metric_ops = None
639
640
scaffold = None
640
- if mode == tf . estimator .ModeKeys .EVAL :
641
+ if mode == tf_estimator .ModeKeys .EVAL :
641
642
class_agnostic = (
642
643
fields .DetectionResultFields .detection_classes not in detections )
643
644
groundtruth = _prepare_groundtruth_for_eval (
@@ -711,8 +712,8 @@ def tpu_scaffold():
711
712
scaffold = tf .train .Scaffold (saver = saver )
712
713
713
714
# EVAL executes on CPU, so use regular non-TPU EstimatorSpec.
714
- if use_tpu and mode != tf . estimator .ModeKeys .EVAL :
715
- return tf . estimator .tpu .TPUEstimatorSpec (
715
+ if use_tpu and mode != tf_estimator .ModeKeys .EVAL :
716
+ return tf_estimator .tpu .TPUEstimatorSpec (
716
717
mode = mode ,
717
718
scaffold_fn = scaffold_fn ,
718
719
predictions = detections ,
@@ -730,7 +731,7 @@ def tpu_scaffold():
730
731
save_relative_paths = True )
731
732
tf .add_to_collection (tf .GraphKeys .SAVERS , saver )
732
733
scaffold = tf .train .Scaffold (saver = saver )
733
- return tf . estimator .EstimatorSpec (
734
+ return tf_estimator .EstimatorSpec (
734
735
mode = mode ,
735
736
predictions = detections ,
736
737
loss = total_loss ,
@@ -895,7 +896,7 @@ def create_estimator_and_inputs(run_config,
895
896
model_fn = model_fn_creator (detection_model_fn , configs , hparams , use_tpu ,
896
897
postprocess_on_cpu )
897
898
if use_tpu_estimator :
898
- estimator = tf . estimator .tpu .TPUEstimator (
899
+ estimator = tf_estimator .tpu .TPUEstimator (
899
900
model_fn = model_fn ,
900
901
train_batch_size = train_config .batch_size ,
901
902
# For each core, only batch size 1 is supported for eval.
@@ -906,7 +907,7 @@ def create_estimator_and_inputs(run_config,
906
907
eval_on_tpu = False , # Eval runs on CPU, so disable eval on TPU
907
908
params = params if params else {})
908
909
else :
909
- estimator = tf . estimator .Estimator (model_fn = model_fn , config = run_config )
910
+ estimator = tf_estimator .Estimator (model_fn = model_fn , config = run_config )
910
911
911
912
# Write the as-run pipeline config to disk.
912
913
if run_config .is_chief and save_final_config :
@@ -951,7 +952,7 @@ def create_train_and_eval_specs(train_input_fn,
951
952
True, the last `EvalSpec` in the list will correspond to training data. The
952
953
rest EvalSpecs in the list are evaluation datas.
953
954
"""
954
- train_spec = tf . estimator .TrainSpec (
955
+ train_spec = tf_estimator .TrainSpec (
955
956
input_fn = train_input_fn , max_steps = train_steps )
956
957
957
958
if eval_spec_names is None :
@@ -966,18 +967,18 @@ def create_train_and_eval_specs(train_input_fn,
966
967
exporter_name = final_exporter_name
967
968
else :
968
969
exporter_name = '{}_{}' .format (final_exporter_name , eval_spec_name )
969
- exporter = tf . estimator .FinalExporter (
970
+ exporter = tf_estimator .FinalExporter (
970
971
name = exporter_name , serving_input_receiver_fn = predict_input_fn )
971
972
eval_specs .append (
972
- tf . estimator .EvalSpec (
973
+ tf_estimator .EvalSpec (
973
974
name = eval_spec_name ,
974
975
input_fn = eval_input_fn ,
975
976
steps = None ,
976
977
exporters = exporter ))
977
978
978
979
if eval_on_train_data :
979
980
eval_specs .append (
980
- tf . estimator .EvalSpec (
981
+ tf_estimator .EvalSpec (
981
982
name = 'eval_on_train' , input_fn = eval_on_train_input_fn , steps = None ))
982
983
983
984
return train_spec , eval_specs
0 commit comments