Repository: pytorch/vision
Pull request: remove unnecessary checks from pad_image_tensor
Closes #6882. As discussed offline, since our implementation already only converts the dtype for non-constant padding

https://github.com/pytorch/vision/blob/1921613a11f0b382be046e77e17281ec60de7fa9/torchvision/transforms/functional_tensor.py#L426-L431

there is not much performance to gain on our side. As explained in https://github.com/pytorch/vision/issues/6818#issuecomment-1292036288, the two possible optimization vectors are edge cases and have to happen in PyTorch core and thus would affect v1 as well as v2. 

Thus, this PR is mostly refactoring.

```
[--------------------------------- pad_image_tensor refactor ----------------------------------]
                                      |     v2 (main)     |   v2 (perf-pad)   |         v1      
1 threads: -------------------------------------------------------------------------------------
      (3, 512, 512), uint8, cpu       |   308 (+-  9) us  |   309 (+-  8) us  |   305 (+-  2) us
      (3, 512, 512), float32, cpu     |   115 (+-  1) us  |   117 (+-  2) us  |   118 (+-  2) us
      (5, 3, 512, 512), uint8, cpu    |  1469 (+- 38) us  |  1476 (+-  7) us  |  1475 (+-  9) us
      (5, 3, 512, 512), float32, cpu  |   910 (+- 89) us  |   893 (+- 64) us  |   919 (+- 60) us

Times are in microseconds (us).
```

cc @vfdev-5 @datumbox @bjuncek

Problem statement:
Optimize `_FT.pad()` on Transforms V2
### 🚀 The feature

The `pad()` implementation for Tensors on stable contains excessive assertions and workarounds to circumvent past limitations of torch pad (unsqueezing for 3d input, float casting for certain modes etc). It's worth reviewing it's implementation and cleaning it up on V2.

Improving its speed might give performance improvements on several key Transforms such as Pad, RandomCrop, etc

cc @vfdev-5 @bjuncek @pmeier

Hints from issue discussion:

Commit message: Merge branch 'main' into perf/pad

Patch to review:
diff --git a/torchvision/prototype/transforms/functional/_geometry.py b/torchvision/prototype/transforms/functional/_geometry.py
index d3e1d94c96c..8ac744b667b 100644
--- a/torchvision/prototype/transforms/functional/_geometry.py
+++ b/torchvision/prototype/transforms/functional/_geometry.py
@@ -4,6 +4,7 @@
 
 import PIL.Image
 import torch
+from torch.nn.functional import pad as torch_pad
 from torchvision.prototype import features
 from torchvision.transforms import functional_pil as _FP, functional_tensor as _FT
 from torchvision.transforms.functional import (
@@ -14,7 +15,6 @@
     pil_to_tensor,
     to_pil_image,
 )
-from torchvision.transforms.functional_tensor import _parse_pad_padding
 
 from ._meta import convert_format_bounding_box, get_spatial_size_image_pil
 
@@ -645,7 +645,32 @@ def rotate(
         return rotate_image_pil(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
 
 
-pad_image_pil = _FP.pad
+def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]:
+    if not isinstance(padding, (int, tuple, list)):
+        raise TypeError(f"`padding` should be an integer or tuple or list of integers, but got {padding}")
+
+    if isinstance(padding, int):
+        if torch.jit.is_scripting():
+            # This maybe unreachable
+            raise ValueError("padding can't be an int while torchscripting, set it as a list [value, ]")
+        pad_left = pad_right = pad_top = pad_bottom = padding
+    else:
+        if len(padding) == 1:
+            pad_left = pad_right = pad_top = pad_bottom = padding[0]
+        elif len(padding) == 2:
+            pad_left = pad_right = padding[0]
+            pad_top = pad_bottom = padding[1]
+        elif len(padding) == 4:
+            pad_left = padding[0]
+            pad_top = padding[1]
+            pad_right = padding[2]
+            pad_bottom = padding[3]
+        else:
+            raise ValueError(
+                f"Padding must be an int or a 1, 2, or 4 element tuple, not a {len(padding)} element tuple"
+            )
+
+    return [pad_left, pad_right, pad_top, pad_bottom]
 
 
 def pad_image_tensor(
@@ -654,50 +679,86 @@ def pad_image_tensor(
     fill: features.FillTypeJIT = None,
     padding_mode: str = "constant",
 ) -> torch.Tensor:
+    # Be aware that while `padding` has order `[left, top, right, bottom]` has order, `torch_padding` uses
+    # `[left, right, top, bottom]`. This stems from the fact that we align our API with PIL, but need to use `torch_pad`
+    # internally.
+    torch_padding = _parse_pad_padding(padding)
+
+    if padding_mode not in ["constant", "edge", "reflect", "symmetric"]:
+        raise ValueError(
+            f"`padding_mode` should be either `'constant'`, `'edge'`, `'reflect'` or `'symmetric'`, "
+            f"but got `'{padding_mode}'`."
+        )
+
     if fill is None:
-        # This is a JIT workaround
-        return _pad_with_scalar_fill(image, padding, fill=None, padding_mode=padding_mode)
-    elif isinstance(fill, (int, float)) or len(fill) == 1:
-        fill_number = fill[0] if isinstance(fill, list) else fill
-        return _pad_with_scalar_fill(image, padding, fill=fill_number, padding_mode=padding_mode)
+        fill = 0
+
+    if isinstance(fill, (int, float)):
+        return _pad_with_scalar_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
+    elif len(fill) == 1:
+        return _pad_with_scalar_fill(image, torch_padding, fill=fill[0], padding_mode=padding_mode)
     else:
-        return _pad_with_vector_fill(image, padding, fill=fill, padding_mode=padding_mode)
+        return _pad_with_vector_fill(image, torch_padding, fill=fill, padding_mode=padding_mode)
 
 
 def _pad_with_scalar_fill(
     image: torch.Tensor,
-    padding: Union[int, List[int]],
-    fill: Union[int, float, None],
-    padding_mode: str = "constant",
+    torch_padding: List[int],
+    fill: Union[int, float],
+    padding_mode: str,
 ) -> torch.Tensor:
     shape = image.shape
     num_channels, height, width = shape[-3:]
 
     if image.numel() > 0:
-        image = _FT.pad(
-            img=image.reshape(-1, num_channels, height, width), padding=padding, fill=fill, padding_mode=padding_mode
-        )
+        image = image.reshape(-1, num_channels, height, width)
+
+        if padding_mode == "edge":
+            # Similar to the padding order, `torch_pad`'s PIL's padding modes don't have the same names. Thus, we map
+            # the PIL name for the padding mode, which we are also using for our API, to the corresponding `torch_pad`
+            # name.
+            padding_mode = "replicate"
+
+        if padding_mode == "constant":
+            image = torch_pad(image, torch_padding, mode=padding_mode, value=float(fill))
+        elif padding_mode in ("reflect", "replicate"):
+            # `torch_pad` only supports `"reflect"` or `"replicate"` padding for floating point inputs.
+            # TODO: See https://github.com/pytorch/pytorch/issues/40763
+            dtype = image.dtype
+            if not image.is_floating_point():
+                needs_cast = True
+                image = image.to(torch.float32)
+            else:
+                needs_cast = False
+
+            image = torch_pad(image, torch_padding, mode=padding_mode)
+
+            if needs_cast:
+                image = image.to(dtype)
+        else:  # padding_mode == "symmetric"
+            image = _FT._pad_symmetric(image, torch_padding)
+
         new_height, new_width = image.shape[-2:]
     else:
-        left, right, top, bottom = _FT._parse_pad_padding(padding)
+        left, right, top, bottom = torch_padding
         new_height = height + top + bottom
         new_width = width + left + right
 
     return image.reshape(shape[:-3] + (num_channels, new_height, new_width))
 
 
-# TODO: This should be removed once pytorch pad supports non-scalar padding values
+# TODO: This should be removed once torch_pad supports non-scalar padding values
 def _pad_with_vector_fill(
     image: torch.Tensor,
-    padding: Union[int, List[int]],
+    torch_padding: List[int],
     fill: List[float],
-    padding_mode: str = "constant",
+    padding_mode: str,
 ) -> torch.Tensor:
     if padding_mode != "constant":
         raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")
 
-    output = _pad_with_scalar_fill(image, padding, fill=0, padding_mode="constant")
-    left, right, top, bottom = _parse_pad_padding(padding)
+    output = _pad_with_scalar_fill(image, torch_padding, fill=0, padding_mode="constant")
+    left, right, top, bottom = torch_padding
     fill = torch.tensor(fill, dtype=image.dtype, device=image.device).reshape(-1, 1, 1)
 
     if top > 0:
@@ -711,6 +772,9 @@ def _pad_with_vector_fill(
     return output
 
 
+pad_image_pil = _FP.pad
+
+
 def pad_mask(
     mask: torch.Tensor,
     padding: Union[int, List[int]],


Code review comments:
diff_hunk: @@ -645,7 +645,32 @@ def rotate(
         return rotate_image_pil(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
 
 
-pad_image_pil = _FP.pad
+def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]:
+    if not isinstance(padding, (int, tuple, list)):
+        raise TypeError(f"`padding` should be an integer or tuple or list of integers, but got {padding}")
+
+    if isinstance(padding, int):
+        if torch.jit.is_scripting():
+            # This maybe unreachable
+            raise ValueError("padding can't be an int while torchscripting, set it as a list [value, ]")
path: torchvision/prototype/transforms/functional/_geometry.py
text: @user1:
I'm probably missing something, we've never used this idiom anywhere else before. Why can't we just remove these lines and let the follow link handle it? Do we get JIT errors?

@author:
TBH, I copy-pasted this. Let me give it a try.

@author:
This error was actually harmful. We had some xfails in place that triggered on that rather then on a JIT failure. So we actually do support integer paddings while scripting.

diff_hunk: @@ -645,7 +645,32 @@ def rotate(
         return rotate_image_pil(inpt, angle, interpolation=interpolation, expand=expand, fill=fill, center=center)
 
 
-pad_image_pil = _FP.pad
+def _parse_pad_padding(padding: Union[int, List[int]]) -> List[int]:
+    if not isinstance(padding, (int, tuple, list)):
+        raise TypeError(f"`padding` should be an integer or tuple or list of integers, but got {padding}")
path: torchvision/prototype/transforms/functional/_geometry.py
text: @user1:
Nit: Shall we move to the if block below where we can first check for int, then tuple/list, else error?