use DSConv2D in heads

official/vision/beta/projects/panoptic_maskrcnn/configs/panoptic_deeplab.py

@@ -200,10 +200,10 @@ def panoptic_deeplab_coco() -> cfg.ExperimentConfig:
                       dilation_rates=aspp_dilation_rates)),
               semantic_head=SemanticHead(
                   level=level,
-                  num_convs=2,
+                  num_convs=1,
                   num_filters=256,
                   kernel_size=5,
-                  use_depthwise_convolution=False,
+                  use_depthwise_convolution=True,
                   upsample_factor=1,
                   low_level=(3, 2),
                   low_level_num_filters=(64, 32),
@@ -211,10 +211,10 @@ def panoptic_deeplab_coco() -> cfg.ExperimentConfig:
                   prediction_kernel_size=1),
               instance_head=InstanceHead(
                   level=level,
-                  num_convs=2,
+                  num_convs=1,
                   num_filters=32,
                   kernel_size=5,
-                  use_depthwise_convolution=False,
+                  use_depthwise_convolution=True,
                   upsample_factor=1,
                   low_level=(3, 2),
                   low_level_num_filters=(32, 16),

official/vision/beta/projects/panoptic_maskrcnn/modeling/heads/panoptic_deeplab_heads.py

@@ -19,7 +19,7 @@ import tensorflow as tf
 
 from official.modeling import tf_utils
 from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import fusion_layers
-from official.vision.beta.ops import spatial_transform_ops
+from official.vision.ops import spatial_transform_ops
 
 
 class PanopticDeeplabHead(tf.keras.layers.Layer):
@@ -109,7 +109,7 @@ class PanopticDeeplabHead(tf.keras.layers.Layer):
     conv_kwargs = {
         'kernel_size': kernel_size if not use_depthwise_convolution else 1,
         'padding': 'same',
-        'use_bias': False,
+        'use_bias': True,
         'kernel_initializer': random_initializer,
         'kernel_regularizer': self._config_dict['kernel_regularizer'],
     }
@@ -127,6 +127,7 @@ class PanopticDeeplabHead(tf.keras.layers.Layer):
         low_level=self._config_dict['low_level'],
         num_projection_filters=self._config_dict['low_level_num_filters'],
         num_output_filters=self._config_dict['fusion_num_output_filters'],
+        use_depthwise_convolution=self._config_dict['use_depthwise_convolution'],
         activation=self._config_dict['activation'],
         use_sync_bn=self._config_dict['use_sync_bn'],
         norm_momentum=self._config_dict['norm_momentum'],
@@ -142,9 +143,9 @@ class PanopticDeeplabHead(tf.keras.layers.Layer):
         self._convs.append(
             tf.keras.layers.DepthwiseConv2D(
                 name='panoptic_deeplab_head_depthwise_conv_{}'.format(i),
-                kernel_size=3,
+                kernel_size=kernel_size,
                 padding='same',
-                use_bias=False,
+                use_bias=True,
                 depthwise_initializer=random_initializer,
                 depthwise_regularizer=self._config_dict['kernel_regularizer'],
                 depth_multiplier=1))

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/fusion_layers.py

@@ -39,6 +39,7 @@ class PanopticDeepLabFusion(tf.keras.layers.Layer):
       low_level: List[int] = [3, 2],
       num_projection_filters: List[int] = [64, 32],
       num_output_filters: int = 256,
+      use_depthwise_convolution: bool = False,
       activation: str = 'relu',
       use_sync_bn: bool = False,
       norm_momentum: float = 0.99,
@@ -52,7 +53,11 @@ class PanopticDeepLabFusion(tf.keras.layers.Layer):
     Args:
       level: An `int` level at which the decoder was appled at.
       low_level: A list of `int` of minimum level to use in feature fusion.
-      num_filters: An `int` number of filters in conv2d layers.
+      num_projection_filters: A list of `int` with number of filters for
+        projection conv2d layers.
+      num_output_filters: An `int` number of filters in output conv2d layers.
+      use_depthwise_convolution: A bool to specify if use depthwise separable
+        convolutions.      
       activation: A `str` name of the activation function.
       use_sync_bn: A `bool` that indicates whether to use synchronized batch
         normalization across different replicas.
@@ -75,6 +80,7 @@ class PanopticDeepLabFusion(tf.keras.layers.Layer):
         'low_level': low_level,
         'num_projection_filters': num_projection_filters,
         'num_output_filters': num_output_filters,
+        'use_depthwise_convolution': use_depthwise_convolution,
         'activation': activation,
         'use_sync_bn': use_sync_bn,
         'norm_momentum': norm_momentum,
@@ -93,7 +99,7 @@ class PanopticDeepLabFusion(tf.keras.layers.Layer):
     conv_op = tf.keras.layers.Conv2D
     conv_kwargs = {
         'padding': 'same',
-        'use_bias': False,
+        'use_bias': True,
         'kernel_initializer': tf.initializers.VarianceScaling(),
         'kernel_regularizer': self._config_dict['kernel_regularizer'],
     }
@@ -116,11 +122,27 @@ class PanopticDeepLabFusion(tf.keras.layers.Layer):
               filters=self._config_dict['num_projection_filters'][i],
               kernel_size=1,
               **conv_kwargs))
-      self._fusion_convs.append(
+      if self._config_dict['use_depthwise_convolution']:
+        depthwise_initializer = tf.keras.initializers.RandomNormal(stddev=0.01)
+        fusion_conv = tf.keras.Sequential([
+            tf.keras.layers.DepthwiseConv2D(
+                kernel_size=5,
+                padding='same',
+                use_bias=True,
+                depthwise_initializer=depthwise_initializer,
+                depthwise_regularizer=self._config_dict['kernel_regularizer'],
+                depth_multiplier=1),
+            bn_op(**bn_kwargs),
             conv_op(
+                filters=self._config_dict['num_output_filters'],
+                kernel_size=1,
+                **conv_kwargs)])
+      else:
+        fusion_conv = conv_op(
             filters=self._config_dict['num_output_filters'],
             kernel_size=5,
-              **conv_kwargs))
+            **conv_kwargs)
+      self._fusion_convs.append(fusion_conv)
       self._projection_norms.append(bn_op(**bn_kwargs))
       self._fusion_norms.append(bn_op(**bn_kwargs))