Internal change

PiperOrigin-RevId: 288388231

Author: yeqingli

official/modeling/training/distributed_executor.py

@@ -248,9 +248,10 @@ def train_step(iterator, num_steps):
             _replicated_step, args=(next(iterator),))
 
       # For reporting, we returns the mean of losses.
-      loss = strategy.reduce(
-          tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None)
-      return loss
+      losses = tf.nest.map_structure(
+          lambda x: strategy.reduce(tf.distribute.ReduceOp.MEAN, x, axis=None),
+          per_replica_losses)
+      return losses
 
     return train_step
 

official/vision/detection/configs/maskrcnn_config.py

@@ -71,6 +71,9 @@
         'min_level': 2,
         'max_level': 6,
         'anchors_per_location': 3,
+        'num_convs': 2,
+        'num_filters': 256,
+        'use_separable_conv': False,
         'use_batch_norm': False,
         'batch_norm': {
             'batch_norm_momentum': 0.997,
@@ -83,7 +86,11 @@
         # Note that `num_classes` is the total number of classes including
         # one background classes whose index is 0.
         'num_classes': 91,
-        'fast_rcnn_mlp_head_dim': 1024,
+        'num_convs': 0,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'num_fcs': 2,
+        'fc_dims': 1024,
         'use_batch_norm': False,
         'batch_norm': {
             'batch_norm_momentum': 0.997,
@@ -95,6 +102,9 @@
     'mrcnn_head': {
         'num_classes': 91,
         'mask_target_size': 28,
+        'num_convs': 4,
+        'num_filters': 256,
+        'use_separable_conv': False,
         'use_batch_norm': False,
         'batch_norm': {
             'batch_norm_momentum': 0.997,

official/vision/detection/dataloader/maskrcnn_parser.py

@@ -353,7 +353,9 @@ def _parse_predict_data(self, data):
         self._anchor_size,
         (image_height, image_width))
 
-    labels = {}
+    labels = {
+        'image_info': image_info,
+    }
 
     if self._mode == ModeKeys.PREDICT_WITH_GT:
       # Converts boxes from normalized coordinates to pixel coordinates.

official/vision/detection/executor/detection_executor.py

@@ -82,7 +82,7 @@ def _replicated_step(inputs):
 
       grads = tape.gradient(loss, trainable_variables)
       optimizer.apply_gradients(zip(grads, trainable_variables))
-      return loss
+      return losses
 
     return _replicated_step
 

official/vision/detection/modeling/architecture/factory.py

@@ -94,14 +94,23 @@ def rpn_head_generator(params):
   return heads.RpnHead(params.min_level,
                        params.max_level,
                        params.anchors_per_location,
+                       params.num_convs,
+                       params.num_filters,
+                       params.use_separable_conv,
+                       params.use_batch_norm,
                        batch_norm_relu=batch_norm_relu_generator(
                            params.batch_norm))
 
 
 def fast_rcnn_head_generator(params):
   """Generator function for Fast R-CNN head architecture."""
   return heads.FastrcnnHead(params.num_classes,
-                            params.fast_rcnn_mlp_head_dim,
+                            params.num_convs,
+                            params.num_filters,
+                            params.use_separable_conv,
+                            params.num_fcs,
+                            params.fc_dims,
+                            params.use_batch_norm,
                             batch_norm_relu=batch_norm_relu_generator(
                                 params.batch_norm))
 
@@ -110,6 +119,10 @@ def mask_rcnn_head_generator(params):
   """Generator function for Mask R-CNN head architecture."""
   return heads.MaskrcnnHead(params.num_classes,
                             params.mask_target_size,
+                            params.num_convs,
+                            params.num_filters,
+                            params.use_separable_conv,
+                            params.use_batch_norm,
                             batch_norm_relu=batch_norm_relu_generator(
                                 params.batch_norm))
 

official/vision/detection/modeling/architecture/fpn.py

@@ -24,6 +24,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 import tensorflow.compat.v2 as tf
 
 from tensorflow.python.keras import backend
@@ -39,6 +41,7 @@ def __init__(self,
                max_level=7,
                fpn_feat_dims=256,
                use_separable_conv=False,
+               use_batch_norm=True,
                batch_norm_relu=nn_ops.BatchNormRelu):
     """FPN initialization function.
 
@@ -48,60 +51,46 @@ def __init__(self,
       fpn_feat_dims: `int` number of filters in FPN layers.
       use_separable_conv: `bool`, if True use separable convolution for
         convolution in FPN layers.
+      use_batch_norm: 'bool', indicating whether batchnorm layers are added.
       batch_norm_relu: an operation that includes a batch normalization layer
         followed by a relu layer(optional).
     """
     self._min_level = min_level
     self._max_level = max_level
     self._fpn_feat_dims = fpn_feat_dims
+    if use_separable_conv:
+      self._conv2d_op = functools.partial(
+          tf.keras.layers.SeparableConv2D, depth_multiplier=1)
+    else:
+      self._conv2d_op = tf.keras.layers.Conv2D
+    self._use_batch_norm = use_batch_norm
     self._batch_norm_relu = batch_norm_relu
 
     self._batch_norm_relus = {}
     self._lateral_conv2d_op = {}
     self._post_hoc_conv2d_op = {}
     self._coarse_conv2d_op = {}
     for level in range(self._min_level, self._max_level + 1):
-      self._batch_norm_relus[level] = batch_norm_relu(
-          relu=False, name='p%d-bn' % level)
-      if use_separable_conv:
-        self._lateral_conv2d_op[level] = tf.keras.layers.SeparableConv2D(
-            filters=self._fpn_feat_dims,
-            kernel_size=(1, 1),
-            padding='same',
-            depth_multiplier=1,
-            name='l%d' % level)
-        self._post_hoc_conv2d_op[level] = tf.keras.layers.SeparableConv2D(
-            filters=self._fpn_feat_dims,
-            strides=(1, 1),
-            kernel_size=(3, 3),
-            padding='same',
-            depth_multiplier=1,
-            name='post_hoc_d%d' % level)
-        self._coarse_conv2d_op[level] = tf.keras.layers.SeparableConv2D(
-            filters=self._fpn_feat_dims,
-            strides=(2, 2),
-            kernel_size=(3, 3),
-            padding='same',
-            depth_multiplier=1,
-            name='p%d' % level)
-      else:
-        self._lateral_conv2d_op[level] = tf.keras.layers.Conv2D(
-            filters=self._fpn_feat_dims,
-            kernel_size=(1, 1),
-            padding='same',
-            name='l%d' % level)
-        self._post_hoc_conv2d_op[level] = tf.keras.layers.Conv2D(
-            filters=self._fpn_feat_dims,
-            strides=(1, 1),
-            kernel_size=(3, 3),
-            padding='same',
-            name='post_hoc_d%d' % level)
-        self._coarse_conv2d_op[level] = tf.keras.layers.Conv2D(
-            filters=self._fpn_feat_dims,
-            strides=(2, 2),
-            kernel_size=(3, 3),
-            padding='same',
-            name='p%d' % level)
+      if self._use_batch_norm:
+        self._batch_norm_relus[level] = batch_norm_relu(
+            relu=False, name='p%d-bn' % level)
+      self._lateral_conv2d_op[level] = self._conv2d_op(
+          filters=self._fpn_feat_dims,
+          kernel_size=(1, 1),
+          padding='same',
+          name='l%d' % level)
+      self._post_hoc_conv2d_op[level] = self._conv2d_op(
+          filters=self._fpn_feat_dims,
+          strides=(1, 1),
+          kernel_size=(3, 3),
+          padding='same',
+          name='post_hoc_d%d' % level)
+      self._coarse_conv2d_op[level] = self._conv2d_op(
+          filters=self._fpn_feat_dims,
+          strides=(2, 2),
+          kernel_size=(3, 3),
+          padding='same',
+          name='p%d' % level)
 
   def __call__(self, multilevel_features, is_training=None):
     """Returns the FPN features for a given multilevel features.
@@ -117,7 +106,7 @@ def __call__(self, multilevel_features, is_training=None):
       [min_level, min_level + 1, ..., max_level]. The values are corresponding
       FPN features with shape [batch_size, height_l, width_l, fpn_feat_dims].
     """
-    input_levels = multilevel_features.keys()
+    input_levels = list(multilevel_features.keys())
     if min(input_levels) > self._min_level:
       raise ValueError(
           'The minimum backbone level %d should be '%(min(input_levels)) +
@@ -146,8 +135,9 @@ def __call__(self, multilevel_features, is_training=None):
         if level > backbone_max_level + 1:
           feats_in = tf.nn.relu(feats_in)
         feats[level] = self._coarse_conv2d_op[level](feats_in)
-      # Adds batch_norm layer.
-      for level in range(self._min_level, self._max_level + 1):
-        feats[level] = self._batch_norm_relus[level](
-            feats[level], is_training=is_training)
+      if self._use_batch_norm:
+        # Adds batch_norm layer.
+        for level in range(self._min_level, self._max_level + 1):
+          feats[level] = self._batch_norm_relus[level](
+              feats[level], is_training=is_training)
     return feats