Internal change

PiperOrigin-RevId: 413811970

official/projects/volumetric_models/modeling/heads/segmentation_heads_3d.py

@@ -155,6 +155,7 @@ class SegmentationHead3D(tf.keras.layers.Layer):
         - key: A `str` of the level of the multilevel features.
         - values: A `tf.Tensor` of the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
+        The first is backbone endpoints, and the second is decoder endpoints.
     Returns:
       segmentation prediction mask: A `tf.Tensor` of the segmentation mask
         scores predicted from input features.

official/vision/beta/modeling/decoders/aspp.py

@@ -149,7 +149,8 @@ class ASPP(tf.keras.layers.Layer):
     return outputs if self._config_dict['output_tensor'] else {level: outputs}
 
   def get_config(self) -> Mapping[str, Any]:
-    return self._config_dict
+    base_config = super().get_config()
+    return dict(list(base_config.items()) + list(self._config_dict.items()))
 
   @classmethod
   def from_config(cls, config, custom_objects=None):

official/vision/beta/modeling/decoders/aspp_test.py

@@ -74,7 +74,10 @@ class ASPPTest(parameterized.TestCase, tf.test.TestCase):
         dropout_rate=0.2,
         use_depthwise_convolution='false',
         spp_layer_version='v1',
-        output_tensor=False)
+        output_tensor=False,
+        dtype='float32',
+        name='aspp',
+        trainable=True)
     network = aspp.ASPP(**kwargs)
 
     expected_config = dict(kwargs)

official/vision/beta/modeling/decoders/factory_test.py

@@ -133,6 +133,10 @@ class FactoryTest(tf.test.TestCase, parameterized.TestCase):
 
     network_config = network.get_config()
     factory_network_config = factory_network.get_config()
+    # Due to calling `super().get_config()` in aspp layer, everything but the
+    # the name of two layer instances are the same, so we force equal name so it
+    # will not give false alarm.
+    factory_network_config['name'] = network_config['name']
 
     self.assertEqual(network_config, factory_network_config)
 

official/vision/beta/modeling/heads/segmentation_heads.py

@@ -202,7 +202,7 @@ class SegmentationHead(tf.keras.layers.Layer):
         kernel_regularizer=self._config_dict['kernel_regularizer'],
         bias_regularizer=self._config_dict['bias_regularizer'])
 
-    super(SegmentationHead, self).build(input_shape)
+    super().build(input_shape)
 
   def call(self, inputs: Tuple[Union[tf.Tensor, Mapping[str, tf.Tensor]],
                                Union[tf.Tensor, Mapping[str, tf.Tensor]]]):
@@ -220,6 +220,7 @@ class SegmentationHead(tf.keras.layers.Layer):
         - key: A `str` of the level of the multilevel features.
         - values: A `tf.Tensor` of the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
+        The first is backbone endpoints, and the second is decoder endpoints.
     Returns:
       segmentation prediction mask: A `tf.Tensor` of the segmentation mask
         scores predicted from input features.
@@ -261,7 +262,8 @@ class SegmentationHead(tf.keras.layers.Layer):
     return self._classifier(x)
 
   def get_config(self):
-    return self._config_dict
+    base_config = super().get_config()
+    return dict(list(base_config.items()) + list(self._config_dict.items()))
 
   @classmethod
   def from_config(cls, config):