Internal change

PiperOrigin-RevId: 479358049

official/vision/modeling/layers/deeplab.py

@@ -87,8 +87,6 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
     self.use_depthwise_convolution = use_depthwise_convolution
 
   def build(self, input_shape):
-    height = input_shape[1]
-    width = input_shape[2]
     channels = input_shape[3]
 
     self.aspp_layers = []
@@ -171,12 +169,7 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
                 axis=bn_axis,
                 momentum=self.batchnorm_momentum,
                 epsilon=self.batchnorm_epsilon),
-            tf.keras.layers.Activation(self.activation),
-            tf.keras.layers.experimental.preprocessing.Resizing(
-                height,
-                width,
-                interpolation=self.interpolation,
-                dtype=tf.float32)
+            tf.keras.layers.Activation(self.activation)
         ]))
 
     self.aspp_layers.append(pool_sequential)
@@ -201,8 +194,12 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
     if training is None:
       training = tf.keras.backend.learning_phase()
     result = []
-    for layer in self.aspp_layers:
+    for i, layer in enumerate(self.aspp_layers):
       result.append(tf.cast(layer(inputs, training=training), inputs.dtype))
+      if i == len(self.aspp_layers) - 1:
+        input_shape = inputs.get_shape().as_list()
+        height, width = input_shape[1:3]
+        result[-1] = tf.image.resize(result[-1], [height, width])
     result = tf.concat(result, axis=-1)
     result = self.projection(result, training=training)
     return result

official/vision/modeling/layers/nn_layers.py

@@ -1146,8 +1146,6 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
       self._bn_axis = 1
 
   def build(self, input_shape):
-    height = input_shape[1]
-    width = input_shape[2]
     channels = input_shape[3]
 
     self.aspp_layers = []
@@ -1220,9 +1218,6 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
 
     self.aspp_layers.append(pooling + [conv2, norm2])
 
-    self._resizing_layer = tf.keras.layers.Resizing(
-        height, width, interpolation=self._interpolation, dtype=tf.float32)
-
     self._projection = [
         tf.keras.layers.Conv2D(
             filters=self._output_channels,
@@ -1254,7 +1249,9 @@ class SpatialPyramidPooling(tf.keras.layers.Layer):
 
       # Apply resize layer to the end of the last set of layers.
       if i == len(self.aspp_layers) - 1:
-        x = self._resizing_layer(x)
+        input_shape = inputs.get_shape().as_list()
+        height, width = input_shape[1:3]
+        x = tf.image.resize(x, [height, width], self._interpolation)
 
       result.append(tf.cast(x, inputs.dtype))
     x = self._concat_layer(result)

official/vision/tasks/semantic_segmentation.py

@@ -308,7 +308,6 @@ class SemanticSegmentationTask(base_task.Task):
       self.iou_metric.update_state(labels, outputs['logits'])
     if metrics:
       self.process_metrics(metrics, labels, outputs)
-      logs.update({m.name: m.result() for m in metrics})
 
     return logs