Changes with no comments.

PiperOrigin-RevId: 351940059

Author: ziyeqinghan

tensorflow_examples/lite/model_maker/third_party/efficientdet/keras/efficientdet_keras.py

@@ -401,18 +401,8 @@ def __init__(self,
     self.bns = []
     self.grad_checkpoint = grad_checkpoint
     self.feature_only = feature_only
-    if separable_conv:
-      conv2d_layer = functools.partial(
-          tf.keras.layers.SeparableConv2D,
-          depth_multiplier=1,
-          data_format=data_format,
-          pointwise_initializer=tf.initializers.variance_scaling(),
-          depthwise_initializer=tf.initializers.variance_scaling())
-    else:
-      conv2d_layer = functools.partial(
-          tf.keras.layers.Conv2D,
-          data_format=data_format,
-          kernel_initializer=tf.random_normal_initializer(stddev=0.01))
+
+    conv2d_layer = self.conv2d_layer(separable_conv, data_format)
     for i in range(self.repeats):
       # If using SeparableConv2D
       self.conv_ops.append(
@@ -435,12 +425,8 @@ def __init__(self,
             ))
       self.bns.append(bn_per_level)
 
-    self.classes = conv2d_layer(
-        num_classes * num_anchors,
-        kernel_size=3,
-        bias_initializer=tf.constant_initializer(-np.log((1 - 0.01) / 0.01)),
-        padding='same',
-        name='class-predict')
+    self.classes = self.classes_layer(
+        conv2d_layer, num_classes, num_anchors, name='class-predict')
 
   @tf.autograph.experimental.do_not_convert
   def _conv_bn_act(self, image, i, level_id, training):
@@ -476,6 +462,33 @@ def call(self, inputs, training, **kwargs):
 
     return class_outputs
 
+  @classmethod
+  def conv2d_layer(cls, separable_conv, data_format):
+    """Gets the conv2d layer in ClassNet class."""
+    if separable_conv:
+      conv2d_layer = functools.partial(
+          tf.keras.layers.SeparableConv2D,
+          depth_multiplier=1,
+          data_format=data_format,
+          pointwise_initializer=tf.initializers.variance_scaling(),
+          depthwise_initializer=tf.initializers.variance_scaling())
+    else:
+      conv2d_layer = functools.partial(
+          tf.keras.layers.Conv2D,
+          data_format=data_format,
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01))
+    return conv2d_layer
+
+  @classmethod
+  def classes_layer(cls, conv2d_layer, num_classes, num_anchors, name):
+    """Gets the classes layer in ClassNet class."""
+    return conv2d_layer(
+        num_classes * num_anchors,
+        kernel_size=3,
+        bias_initializer=tf.constant_initializer(-np.log((1 - 0.01) / 0.01)),
+        padding='same',
+        name=name)
+
 
 class BoxNet(tf.keras.layers.Layer):
   """Box regression network."""
@@ -574,28 +587,8 @@ def __init__(self,
                 name='box-%d-bn-%d' % (i, level)))
       self.bns.append(bn_per_level)
 
-    if self.separable_conv:
-      self.boxes = tf.keras.layers.SeparableConv2D(
-          filters=4 * self.num_anchors,
-          depth_multiplier=1,
-          pointwise_initializer=tf.initializers.variance_scaling(),
-          depthwise_initializer=tf.initializers.variance_scaling(),
-          data_format=self.data_format,
-          kernel_size=3,
-          activation=None,
-          bias_initializer=tf.zeros_initializer(),
-          padding='same',
-          name='box-predict')
-    else:
-      self.boxes = tf.keras.layers.Conv2D(
-          filters=4 * self.num_anchors,
-          kernel_initializer=tf.random_normal_initializer(stddev=0.01),
-          data_format=self.data_format,
-          kernel_size=3,
-          activation=None,
-          bias_initializer=tf.zeros_initializer(),
-          padding='same',
-          name='box-predict')
+      self.boxes = self.boxes_layer(
+          separable_conv, num_anchors, data_format, name='box-predict')
 
   @tf.autograph.experimental.do_not_convert
   def _conv_bn_act(self, image, i, level_id, training):
@@ -632,6 +625,32 @@ def call(self, inputs, training):
 
     return box_outputs
 
+  @classmethod
+  def boxes_layer(cls, separable_conv, num_anchors, data_format, name):
+    """Gets the conv2d layer in BoxNet class."""
+    if separable_conv:
+      return tf.keras.layers.SeparableConv2D(
+          filters=4 * num_anchors,
+          depth_multiplier=1,
+          pointwise_initializer=tf.initializers.variance_scaling(),
+          depthwise_initializer=tf.initializers.variance_scaling(),
+          data_format=data_format,
+          kernel_size=3,
+          activation=None,
+          bias_initializer=tf.zeros_initializer(),
+          padding='same',
+          name=name)
+    else:
+      return tf.keras.layers.Conv2D(
+          filters=4 * num_anchors,
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01),
+          data_format=data_format,
+          kernel_size=3,
+          activation=None,
+          bias_initializer=tf.zeros_initializer(),
+          padding='same',
+          name=name)
+
 
 class SegmentationHead(tf.keras.layers.Layer):
   """Keras layer for semantic segmentation head."""

tensorflow_examples/lite/model_maker/third_party/efficientdet/keras/train_lib.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 """Training related libraries."""
-import functools
 import math
 import os
 import re
@@ -863,40 +862,19 @@ def __init__(self, config, hub_module_url, name=''):
     # class/box output prediction network.
     num_anchors = len(config.aspect_ratios) * config.num_scales
 
-    if config.separable_conv:
-      conv2d_layer = functools.partial(
-          tf.keras.layers.SeparableConv2D, depth_multiplier=1)
-    else:
-      conv2d_layer = tf.keras.layers.Conv2D
-    self.classes = conv2d_layer(
-        config.num_classes * num_anchors,
-        kernel_size=3,
-        bias_initializer=tf.constant_initializer(-np.log((1 - 0.01) / 0.01)),
-        padding='same',
+    conv2d_layer = efficientdet_keras.ClassNet.conv2d_layer(
+        config.separable_conv, config.data_format)
+    self.classes = efficientdet_keras.ClassNet.classes_layer(
+        conv2d_layer,
+        config.num_classes,
+        num_anchors,
         name='class_net/class-predict')
 
-    if config.separable_conv:
-      self.boxes = tf.keras.layers.SeparableConv2D(
-          filters=4 * num_anchors,
-          depth_multiplier=1,
-          pointwise_initializer=tf.initializers.variance_scaling(),
-          depthwise_initializer=tf.initializers.variance_scaling(),
-          data_format=config.data_format,
-          kernel_size=3,
-          activation=None,
-          bias_initializer=tf.zeros_initializer(),
-          padding='same',
-          name='box_net/box-predict')
-    else:
-      self.boxes = tf.keras.layers.Conv2D(
-          filters=4 * num_anchors,
-          kernel_initializer=tf.random_normal_initializer(stddev=0.01),
-          data_format=config.data_format,
-          kernel_size=3,
-          activation=None,
-          bias_initializer=tf.zeros_initializer(),
-          padding='same',
-          name='box_net/box-predict')
+    self.boxes = efficientdet_keras.BoxNet.boxes_layer(
+        config.separable_conv,
+        num_anchors,
+        config.data_format,
+        name='box_net/box-predict')
 
     log_dir = os.path.join(self.config.model_dir, 'train_images')
     self.summary_writer = tf.summary.create_file_writer(log_dir)