do not rewrite estimator batch norm

Author: fsx950223

efficientdet/hparams_config.py

@@ -145,7 +145,7 @@ def merge_dict_recursive(target, src):
           """Recursively merge two nested dictionary."""
           for k in src.keys():
             if ((k in target and isinstance(target[k], dict) and
-                 isinstance(src[k], collections.abc.Mapping))):
+                 isinstance(src[k], collections.Mapping))):
               merge_dict_recursive(target[k], src[k])
             else:
               target[k] = src[k]

efficientdet/keras/efficientdet_keras.py

@@ -27,10 +27,7 @@
 from keras import postprocess
 from keras import tfmot
 from keras import util_keras
-# pylint: disable=arguments-differ  # fo keras layers.
-utils.BatchNormalization = util_keras.get_batch_norm(tf.keras.layers.BatchNormalization)
-utils.SyncBatchNormalization = util_keras.get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)
-utils.TpuBatchNormalization = util_keras.get_batch_norm(tf.keras.layers.experimental.SyncBatchNormalization)
+
 
 def add_n(nodes):
   """A customized add_n to add up a list of tensors."""

efficientdet/keras/util_keras.py

@@ -175,9 +175,15 @@ def fp16_to_fp32_nested(input_nested):
     return input_nested
   return out_tensor_dict
 
+
 def get_batch_norm(bn_class):
   def _wrapper(*args, **kwargs):
     if not kwargs.get('name', None):
       kwargs['name'] = 'tpu_batch_normalization'
     return bn_class(*args, **kwargs)
-  return _wrapper
+  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)