Allow setting pointwise initializer and regularizer when use_depthwise is true.

PiperOrigin-RevId: 365899795

research/object_detection/builders/box_predictor_builder.py

@@ -329,6 +329,8 @@ def build_weight_shared_convolutional_keras_box_predictor(
     share_prediction_tower=False,
     apply_batch_norm=True,
     use_depthwise=False,
+    apply_conv_hyperparams_to_heads=False,
+    apply_conv_hyperparams_pointwise=False,
     score_converter_fn=tf.identity,
     box_encodings_clip_range=None,
     name='WeightSharedConvolutionalBoxPredictor',
@@ -369,6 +371,14 @@ def build_weight_shared_convolutional_keras_box_predictor(
     apply_batch_norm: Whether to apply batch normalization to conv layers in
       this predictor.
     use_depthwise: Whether to use depthwise separable conv2d instead of conv2d.
+    apply_conv_hyperparams_to_heads: Whether to apply conv_hyperparams to
+      depthwise seperable convolution layers in the box and class heads. By
+      default, the conv_hyperparams are only applied to layers in the predictor
+      tower when using depthwise separable convolutions.
+    apply_conv_hyperparams_pointwise: Whether to apply the conv_hyperparams to
+      the pointwise_initializer and pointwise_regularizer when using depthwise
+      separable convolutions. By default, conv_hyperparams are only applied to
+      the depthwise initializer and regularizer when use_depthwise is true.
     score_converter_fn: Callable score converter to perform elementwise op on
       class scores.
     box_encodings_clip_range: Min and max values for clipping the box_encodings.
@@ -391,6 +401,7 @@ def build_weight_shared_convolutional_keras_box_predictor(
       conv_hyperparams=conv_hyperparams,
       num_predictions_per_location=num_predictions_per_location,
       use_depthwise=use_depthwise,
+      apply_conv_hyperparams_to_heads=apply_conv_hyperparams_to_heads,
       box_encodings_clip_range=box_encodings_clip_range,
       name='WeightSharedConvolutionalBoxHead')
   class_prediction_head = keras_class_head.WeightSharedConvolutionalClassHead(
@@ -403,6 +414,7 @@ def build_weight_shared_convolutional_keras_box_predictor(
       num_predictions_per_location=num_predictions_per_location,
       class_prediction_bias_init=class_prediction_bias_init,
       use_depthwise=use_depthwise,
+      apply_conv_hyperparams_to_heads=apply_conv_hyperparams_to_heads,
       score_converter_fn=score_converter_fn,
       name='WeightSharedConvolutionalClassHead')
   other_heads = {}
@@ -423,6 +435,7 @@ def build_weight_shared_convolutional_keras_box_predictor(
           apply_batch_norm=apply_batch_norm,
           share_prediction_tower=share_prediction_tower,
           use_depthwise=use_depthwise,
+          apply_conv_hyperparams_pointwise=apply_conv_hyperparams_pointwise,
           name=name))
 
 
@@ -920,6 +933,10 @@ def build_keras(hyperparams_fn, freeze_batchnorm, inplace_batchnorm_update,
         share_prediction_tower=config_box_predictor.share_prediction_tower,
         apply_batch_norm=apply_batch_norm,
         use_depthwise=config_box_predictor.use_depthwise,
+        apply_conv_hyperparams_to_heads=(
+            config_box_predictor.apply_conv_hyperparams_to_heads),
+        apply_conv_hyperparams_pointwise=(
+            config_box_predictor.apply_conv_hyperparams_pointwise),
         score_converter_fn=score_converter_fn,
         box_encodings_clip_range=box_encodings_clip_range,
         keyword_args=keyword_args)

research/object_detection/predictors/convolutional_keras_box_predictor.py

@@ -236,6 +236,7 @@ class WeightSharedConvolutionalBoxPredictor(box_predictor.KerasBoxPredictor):
                apply_batch_norm=False,
                share_prediction_tower=False,
                use_depthwise=False,
+               apply_conv_hyperparams_pointwise=False,
                name=None):
     """Constructor.
 
@@ -269,6 +270,10 @@ class WeightSharedConvolutionalBoxPredictor(box_predictor.KerasBoxPredictor):
         prediction head, class prediction head and other heads.
       use_depthwise: Whether to use depthwise separable conv2d instead of
        regular conv2d.
+      apply_conv_hyperparams_pointwise: Whether to apply the conv_hyperparams to
+        the pointwise_initializer and pointwise_regularizer when using depthwise
+        separable convolutions. By default, conv_hyperparams are only applied to
+        the depthwise initializer and regularizer when use_depthwise is true.
       name: A string name scope to assign to the model. If `None`, Keras
         will auto-generate one from the class name.
     """
@@ -294,6 +299,7 @@ class WeightSharedConvolutionalBoxPredictor(box_predictor.KerasBoxPredictor):
     self._apply_batch_norm = apply_batch_norm
     self._share_prediction_tower = share_prediction_tower
     self._use_depthwise = use_depthwise
+    self._apply_conv_hyperparams_pointwise = apply_conv_hyperparams_pointwise
 
     # Additional projection layers to bring all feature maps to uniform
     # channels.
@@ -344,6 +350,9 @@ class WeightSharedConvolutionalBoxPredictor(box_predictor.KerasBoxPredictor):
           # so we remap the kernel_* to depthwise_* here.
           kwargs['depthwise_regularizer'] = kwargs['kernel_regularizer']
           kwargs['depthwise_initializer'] = kwargs['kernel_initializer']
+          if self._apply_conv_hyperparams_pointwise:
+            kwargs['pointwise_regularizer'] = kwargs['kernel_regularizer']
+            kwargs['pointwise_initializer'] = kwargs['kernel_initializer']
           conv_layers.append(
               tf.keras.layers.SeparableConv2D(
                   self._depth, [self._kernel_size, self._kernel_size],

research/object_detection/predictors/heads/keras_box_head.py

@@ -248,6 +248,7 @@ class WeightSharedConvolutionalBoxHead(head.KerasHead):
                conv_hyperparams,
                kernel_size=3,
                use_depthwise=False,
+               apply_conv_hyperparams_to_heads=False,
                box_encodings_clip_range=None,
                return_flat_predictions=True,
                name=None):
@@ -262,6 +263,10 @@ class WeightSharedConvolutionalBoxHead(head.KerasHead):
       kernel_size: Size of final convolution kernel.
       use_depthwise: Whether to use depthwise convolutions for prediction steps.
         Default is False.
+      apply_conv_hyperparams_to_heads: Whether to apply conv_hyperparams to
+        depthwise seperable convolution layers in the box and class heads. By
+        default, the conv_hyperparams are only applied to layers in the
+        predictor tower when using depthwise separable convolutions.
       box_encodings_clip_range: Min and max values for clipping box_encodings.
       return_flat_predictions: If true, returns flattened prediction tensor
         of shape [batch, height * width * num_predictions_per_location,
@@ -282,19 +287,26 @@ class WeightSharedConvolutionalBoxHead(head.KerasHead):
     self._kernel_size = kernel_size
     self._num_predictions_per_location = num_predictions_per_location
     self._use_depthwise = use_depthwise
+    self._apply_conv_hyperparams_to_heads = apply_conv_hyperparams_to_heads
     self._box_encodings_clip_range = box_encodings_clip_range
     self._return_flat_predictions = return_flat_predictions
 
     self._box_encoder_layers = []
 
     if self._use_depthwise:
+      kwargs = conv_hyperparams.params(use_bias=True)
+      if self._apply_conv_hyperparams_to_heads:
+        kwargs['depthwise_regularizer'] = kwargs['kernel_regularizer']
+        kwargs['depthwise_initializer'] = kwargs['kernel_initializer']
+        kwargs['pointwise_regularizer'] = kwargs['kernel_regularizer']
+        kwargs['pointwise_initializer'] = kwargs['kernel_initializer']
       self._box_encoder_layers.append(
           tf.keras.layers.SeparableConv2D(
               num_predictions_per_location * self._box_code_size,
               [self._kernel_size, self._kernel_size],
               padding='SAME',
               name='BoxPredictor',
-              **conv_hyperparams.params(use_bias=True)))
+              **kwargs))
     else:
       self._box_encoder_layers.append(
           tf.keras.layers.Conv2D(

research/object_detection/predictors/heads/keras_class_head.py

@@ -251,6 +251,7 @@ class WeightSharedConvolutionalClassHead(head.KerasHead):
                use_dropout=False,
                dropout_keep_prob=0.8,
                use_depthwise=False,
+               apply_conv_hyperparams_to_heads=False,
                score_converter_fn=tf.identity,
                return_flat_predictions=True,
                name=None):
@@ -270,6 +271,10 @@ class WeightSharedConvolutionalClassHead(head.KerasHead):
       dropout_keep_prob: Probability of keeping activiations.
       use_depthwise: Whether to use depthwise convolutions for prediction
         steps. Default is False.
+      apply_conv_hyperparams_to_heads: Whether to apply conv_hyperparams to
+        depthwise seperable convolution layers in the box and class heads. By
+        default, the conv_hyperparams are only applied to layers in the
+        predictor tower when using depthwise separable convolutions.
       score_converter_fn: Callable elementwise nonlinearity (that takes tensors
         as inputs and returns tensors).
       return_flat_predictions: If true, returns flattened prediction tensor
@@ -294,6 +299,7 @@ class WeightSharedConvolutionalClassHead(head.KerasHead):
     self._use_dropout = use_dropout
     self._dropout_keep_prob = dropout_keep_prob
     self._use_depthwise = use_depthwise
+    self._apply_conv_hyperparams_to_heads = apply_conv_hyperparams_to_heads
     self._score_converter_fn = score_converter_fn
     self._return_flat_predictions = return_flat_predictions
 
@@ -303,6 +309,12 @@ class WeightSharedConvolutionalClassHead(head.KerasHead):
       self._class_predictor_layers.append(
           tf.keras.layers.Dropout(rate=1.0 - self._dropout_keep_prob))
     if self._use_depthwise:
+      kwargs = conv_hyperparams.params(use_bias=True)
+      if self._apply_conv_hyperparams_to_heads:
+        kwargs['depthwise_regularizer'] = kwargs['kernel_regularizer']
+        kwargs['depthwise_initializer'] = kwargs['kernel_initializer']
+        kwargs['pointwise_regularizer'] = kwargs['kernel_regularizer']
+        kwargs['pointwise_initializer'] = kwargs['kernel_initializer']
       self._class_predictor_layers.append(
           tf.keras.layers.SeparableConv2D(
               num_predictions_per_location * self._num_class_slots,
@@ -313,7 +325,7 @@ class WeightSharedConvolutionalClassHead(head.KerasHead):
               name='ClassPredictor',
               bias_initializer=tf.constant_initializer(
                   self._class_prediction_bias_init),
-              **conv_hyperparams.params(use_bias=True)))
+              **kwargs))
     else:
       self._class_predictor_layers.append(
           tf.keras.layers.Conv2D(

research/object_detection/protos/box_predictor.proto

@@ -66,11 +66,23 @@ message ConvolutionalBoxPredictor {
 }
 
 // Configuration proto for weight shared convolutional box predictor.
-// Next id: 19
+// Next id: 21
 message WeightSharedConvolutionalBoxPredictor {
   // Hyperparameters for convolution ops used in the box predictor.
   optional Hyperparams conv_hyperparams = 1;
 
+  // Whether the `conv_hyperparams` should apply to depthwise separable
+  // convolution layers in the box and class heads, in addition to the layers in
+  // the predictor tower. By default, the `conv_hyperparams` are only applied to
+  // layers in the predictor tower when use_depthwise is true.
+  optional bool apply_conv_hyperparams_to_heads = 19 [default = false];
+
+  // Whether the `conv_hyperparams` should apply to the `pointwise_initializer`
+  // and `pointwise_regularizer` when using depthwise separable convolutions in
+  // the prediction tower layers. By default, the `conv_hyperparams` only apply
+  // to the `depthwise_initializer` and `depthwise_regularizer`.
+  optional bool apply_conv_hyperparams_pointwise = 20 [default = false];
+
   // Number of the additional conv layers before the predictor.
   optional int32 num_layers_before_predictor = 4 [default = 0];