fixed same padding in causal convolutions

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

@@ -720,6 +720,7 @@ class CausalConvMixin:
     Returns:
       A list of paddings for `tf.pad`.
     """
+    shape_in = inputs.shape[1:-1]
     del inputs
 
     if tf.keras.backend.image_data_format() == 'channels_first':
@@ -730,7 +731,11 @@ class CausalConvMixin:
          (self.kernel_size[i] - 1) * (self.dilation_rate[i] - 1))
         for i in range(self.rank)
     ]
-    pad_total = [kernel_size_effective[i] - 1 for i in range(self.rank)]
+    pad_total = [max(kernel_size_effective[i] - (self.strides[i]), 0)
+                 if (shape_in[i]%self.strides[i]) == 0 else
+                 max(kernel_size_effective[i] - 
+                     (shape_in[i]%self.strides[i]), 0)
+                 for i in range(self.rank)]
     pad_beg = [pad_total[i] // 2 for i in range(self.rank)]
     pad_end = [pad_total[i] - pad_beg[i] for i in range(self.rank)]
     padding = [[pad_beg[i], pad_end[i]] for i in range(self.rank)]

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

@@ -274,14 +274,14 @@ class NNLayersTest(parameterized.TestCase, tf.test.TestCase):
     predicted = conv3d(padded_inputs)
 
     expected = tf.constant(
-        [[[[[12., 12., 12.],
+        [[[[[27., 27., 27.],
             [18., 18., 18.]],
            [[18., 18., 18.],
-            [27., 27., 27.]]],
-          [[[24., 24., 24.],
+            [12., 12., 12.]]],
+          [[[54., 54., 54.],
             [36., 36., 36.]],
            [[36., 36., 36.],
-            [54., 54., 54.]]]]])
+            [24., 24., 24.]]]]])
 
     self.assertEqual(predicted.shape, expected.shape)
     self.assertAllClose(predicted, expected)
@@ -311,14 +311,14 @@ class NNLayersTest(parameterized.TestCase, tf.test.TestCase):
     predicted = conv3d(padded_inputs)
 
     expected = tf.constant(
-        [[[[[4.0, 4.0, 4.0],
+        [[[[[9.0, 9.0, 9.0],
             [6.0, 6.0, 6.0]],
            [[6.0, 6.0, 6.0],
-            [9.0, 9.0, 9.0]]],
-          [[[8.0, 8.0, 8.0],
+            [4.0, 4.0, 4.0]]],
+          [[[18.0, 18.0, 18.0],
             [12., 12., 12.]],
            [[12., 12., 12.],
-            [18., 18., 18.]]]]])
+            [8., 8., 8.]]]]])
 
     self.assertEqual(predicted.shape, expected.shape)
     self.assertAllClose(predicted, expected)