fix update ops

Author: fsx950223

efficientdet/keras/util_keras.py

@@ -183,7 +183,7 @@ def _wrapper(*args, **kwargs):
     return bn_class(*args, **kwargs)
   return _wrapper
 
-
-utils.BatchNormalization = get_batch_norm(tf.keras.layers.BatchNormalization)
-utils.SyncBatchNormalization = get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)
-utils.TpuBatchNormalization = get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)
+if tf.compat.v1.executing_eagerly_outside_functions():
+  utils.BatchNormalization = get_batch_norm(tf.keras.layers.BatchNormalization)
+  utils.SyncBatchNormalization = get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)
+  utils.TpuBatchNormalization = get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)

efficientdet/main.py

@@ -57,7 +57,7 @@ def _filter_fn(item):
 
 custom_gradient.get_variable_by_name = get_variable_by_name
 import tensorflow.compat.v1 as tf
-
+tf.disable_eager_execution()
 import dataloader
 import det_model_fn
 import hparams_config
@@ -158,7 +158,6 @@ def _filter_fn(item):
 
 def main(_):
   if FLAGS.strategy == 'tpu':
-    tf.disable_eager_execution()
     tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
         FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
     tpu_grpc_url = tpu_cluster_resolver.get_master()