Repository: keras-team/keras

Issue:
Use the `padding` arg when possible for torch backend
To support `padding="same"`, we implemented a custom padding by calling `torch.nn.functional.pad` for ops including `conv`, `max_pool`, and `average_pool`.

However, it takes more memory and slows down the model.

Torch does support some cases of `padding="same"`.
We should only use this manuall padding method when the native ops in torch doesn't support it.



Discussion:

Resolving patch:
diff --git a/keras/backend/torch/nn.py b/keras/backend/torch/nn.py
--- a/keras/backend/torch/nn.py
+++ b/keras/backend/torch/nn.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 import torch.nn.functional as tnn
 import tree
@@ -112,20 +111,30 @@ def _compute_padding_length(
     input_length, kernel_length, stride, dilation_rate=1
 ):
     """Compute padding length along one dimension."""
-    if (input_length - 1) % stride == 0:
-        total_padding_length = dilation_rate * (kernel_length - 1)
-    else:
-        total_padding_length = (
-            dilation_rate * (kernel_length - 1) - (input_length - 1) % stride
-        )
-    left_padding = int(np.floor(total_padding_length / 2))
-    right_padding = int(np.ceil(total_padding_length / 2))
+    total_padding_length = (
+        dilation_rate * (kernel_length - 1) - (input_length - 1) % stride
+    )
+    left_padding = total_padding_length // 2
+    right_padding = (total_padding_length + 1) // 2
     return (left_padding, right_padding)
 
 
 def _apply_same_padding(
     inputs, kernel_size, strides, operation_type, dilation_rate=1
 ):
+    """Apply same padding to the input tensor.
+
+    This function will evaluate if the padding value is compatible with torch
+    functions. To avoid calling `pad()` as much as possible, which may cause
+    performance or memory issues, when compatible, it does not apply the padding
+    to the tensor, but returns the input tensor and the padding value to pass to
+    the torch functions. If not compatible, it returns the padded tensor and 0
+    as the padding value.
+
+    Returns:
+        tensor: A padded tensor or the inputs.
+        padding: The padding value, ready to pass to the torch functions.
+    """
     spatial_shape = inputs.shape[2:]
     num_spatial_dims = len(spatial_shape)
     padding = ()
@@ -144,9 +153,15 @@ def _apply_same_padding(
                 spatial_shape[i], kernel_size[i], strides[i], dilation_rate[i]
             )
             mode = "constant"
-        padding = padding_size + padding
+        padding = (padding_size,) + padding
 
-    return tnn.pad(inputs, padding, mode=mode)
+    if all([left == right for left, right in padding]):
+        return inputs, [left for left, _ in padding]
+
+    flattened_padding = tuple(
+        value for left_and_right in padding for value in left_and_right
+    )
+    return tnn.pad(inputs, pad=flattened_padding, mode=mode), 0
 
 
 def _transpose_spatial_inputs(inputs):
@@ -215,9 +230,11 @@ def max_pool(
     if padding == "same":
         # Torch does not natively support `"same"` padding, we need to manually
         # apply the right amount of padding to `inputs`.
-        inputs = _apply_same_padding(
+        inputs, padding = _apply_same_padding(
             inputs, pool_size, strides, operation_type="pooling"
         )
+    else:
+        padding = 0
 
     device = get_device()
     # Torch max pooling ops do not support symbolic tensors.
@@ -228,11 +245,17 @@ def max_pool(
         )
 
     if num_spatial_dims == 1:
-        outputs = tnn.max_pool1d(inputs, kernel_size=pool_size, stride=strides)
+        outputs = tnn.max_pool1d(
+            inputs, kernel_size=pool_size, stride=strides, padding=padding
+        )
     elif num_spatial_dims == 2:
-        outputs = tnn.max_pool2d(inputs, kernel_size=pool_size, stride=strides)
+        outputs = tnn.max_pool2d(
+            inputs, kernel_size=pool_size, stride=strides, padding=padding
+        )
     elif num_spatial_dims == 3:
-        outputs = tnn.max_pool3d(inputs, kernel_size=pool_size, stride=strides)
+        outputs = tnn.max_pool3d(
+            inputs, kernel_size=pool_size, stride=strides, padding=padding
+        )
     else:
         raise ValueError(
             "Inputs to pooling op must have ndim=3, 4 or 5, "
@@ -283,7 +306,9 @@ def average_pool(
                 # Handle unequal padding.
                 # `torch.nn.pad` sets padding value in the reverse order.
                 uneven_padding = [0, 1] + uneven_padding
-        inputs = tnn.pad(inputs, uneven_padding)
+        # Only call tnn.pad when needed.
+        if len(uneven_padding) > 0:
+            inputs = tnn.pad(inputs, uneven_padding)
 
     if num_spatial_dims == 1:
         outputs = tnn.avg_pool1d(
@@ -341,14 +366,13 @@ def conv(
     if padding == "same" and any(d != 1 for d in tree.flatten(strides)):
         # Torch does not support this case in conv2d().
         # Manually pad the tensor.
-        inputs = _apply_same_padding(
+        inputs, padding = _apply_same_padding(
             inputs,
             kernel.shape[2:],
             strides,
             operation_type="conv",
             dilation_rate=dilation_rate,
         )
-        padding = 0
     channels = inputs.shape[1]
     kernel_in_channels = kernel.shape[1]
     if channels % kernel_in_channels > 0:


Test patch: