Revert causal padding until new checkpoints are trained.

PiperOrigin-RevId: 382606601

official/vision/beta/modeling/layers/nn_layers.py

@@ -26,6 +26,10 @@ from official.modeling import tf_utils
 States = Dict[str, tf.Tensor]
 Activation = Union[str, Callable]
 
+# TODO(dankondratyuk): keep legacy padding until new checkpoints are trained.
+# Otherwise, accuracy will be affected.
+LEGACY_PADDING = True
+
 
 def make_divisible(value: float,
                    divisor: int,
@@ -725,7 +729,8 @@ class CausalConvMixin:
   def _compute_buffered_causal_padding(self,
                                        inputs: tf.Tensor,
                                        use_buffered_input: bool = False,
-                                       time_axis: int = 1) -> List[List[int]]:
+                                       time_axis: int = 1,
+                                       ) -> List[List[int]]:
     """Calculates padding for 'causal' option for conv layers.
 
     Args:
@@ -747,6 +752,10 @@ class CausalConvMixin:
          (self.kernel_size[i] - 1) * (self.dilation_rate[i] - 1))
         for i in range(self.rank)
     ]
+    if LEGACY_PADDING:
+      # Apply legacy padding that does not take into account spatial strides
+      pad_total = [kernel_size_effective[i] - 1 for i in range(self.rank)]
+    else:
       pad_total = [kernel_size_effective[0] - 1]
       for i in range(1, self.rank):
         overlap = (input_shape[i] - 1) % self.strides[i] + 1

official/vision/beta/modeling/layers/nn_layers_test.py

@@ -24,6 +24,10 @@ from official.vision.beta.modeling.layers import nn_layers
 
 class NNLayersTest(parameterized.TestCase, tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    nn_layers.LEGACY_PADDING = False
+
   def test_hard_swish(self):
     activation = tf.keras.layers.Activation('hard_swish')
     output = activation(tf.constant([-3, -1.5, 0, 3]))