Add scope reuse support.

PiperOrigin-RevId: 415459931

Author: tensorflower-gardener

tensorflow_examples/lite/model_maker/third_party/efficientdet/efficientdet_arch.py

@@ -440,10 +440,16 @@ def fuse_features(nodes, weight_method):
     nodes = tf.stack(nodes, axis=-1)
     new_node = tf.reduce_sum(nodes * normalized_weights, -1)
   elif weight_method == 'fastattn':
-    edge_weights = [
-        tf.nn.relu(tf.cast(tf.Variable(1.0, name='WSM'), dtype=dtype))
-        for _ in nodes
-    ]
+    edge_weights = []
+    for i, _ in enumerate(nodes):
+      edge_weights.append(
+          tf.nn.relu(
+              tf.get_variable(
+                  'WSM' if i == 0 else f'WSM_{i}',
+                  shape=[],
+                  initializer=tf.ones_initializer(),
+                  dtype=dtype)))
+
     weights_sum = tf.add_n(edge_weights)
     nodes = [nodes[i] * edge_weights[i] / (weights_sum + 0.0001)
              for i in range(len(nodes))]

tensorflow_examples/lite/model_maker/third_party/efficientdet/utils.py

@@ -17,6 +17,7 @@
 import os
 from typing import Text, Tuple, Union
 from absl import logging
+import gin
 import numpy as np
 import tensorflow.compat.v1 as tf
 import tensorflow.compat.v2 as tf2
@@ -266,10 +267,43 @@ def batch_norm_class(is_training, strategy=None):
   else:
     return BatchNormalization
 
+# A cache of variable scope to BatchNorm layer.
+_BN_LAYER_CACHE = {}
 
-def batch_normalization(inputs, training=False, strategy=None, **kwargs):
-  """A wrapper for TpuBatchNormalization."""
-  bn_layer = batch_norm_class(training, strategy)(**kwargs)
+
+@gin.configurable
+def batch_normalization(inputs,
+                        training=False,
+                        strategy=None,
+                        reuse_scope: bool = False,
+                        **kwargs):
+  """A wrapper for TpuBatchNormalization.
+
+  Keras layers are incompatible with automatic tf scope reuse.
+
+  Supports reuse of the exiting variable scope when a model is called multiple
+  times. Otherwise, checkpoint weights would not be restored correctly.
+
+  Args:
+    inputs: Input to BatchNorm layer.
+    training: Argument of Keras BatchNorm layer.
+    strategy: Argument of Keras BatchNorm layer.
+    reuse_scope: Whether to reuse existing layer in same scope.
+    **kwargs: Arguments passed to Keras BatchNorm layer.
+
+  Returns:
+    Result of BatchNorm applied to inputs.
+  """
+  if reuse_scope:
+    scope_name = tf.get_variable_scope().name
+    if scope_name in _BN_LAYER_CACHE:
+      bn_layer = _BN_LAYER_CACHE[scope_name]
+      logging.info('Reusing variable scope %s', scope_name)
+    else:
+      bn_layer = batch_norm_class(training, strategy)(**kwargs)
+      _BN_LAYER_CACHE[scope_name] = bn_layer
+  else:
+    bn_layer = batch_norm_class(training, strategy)(**kwargs)
   return bn_layer(inputs, training=training)