Added symmetric padding mode for Tensors

[vfdev-5]

Fixes #2350

Description:

• Added code to support symmetric padding mode for Tensors
• Added test of pad from functional_tensor

This PR is related to #2371 which updates docs and adds more tests with padding mode for class transforms. It is possible to fetch the tests code from that PR and update docs here to the latest state of the API.

Important: Current implementation uses advanced indexing to perform padding. Due to jit limitations (pytorch/pytorch#34837), we can only support 3D and 4D (maybe we should add 2D) tensors.

[codecov]

Codecov Report

Merging #2373 into master will decrease coverage by 0.10%.
The diff coverage is 90.47%.

@@            Coverage Diff             @@
##           master    #2373      +/-   ##
==========================================
- Coverage   68.59%   68.49%   -0.11%     
==========================================
  Files          93       94       +1     
  Lines        7713     7770      +57     
  Branches     1196     1214      +18     
==========================================
+ Hits         5291     5322      +31     
- Misses       2077     2098      +21     
- Partials      345      350       +5     

Impacted Files	Coverage Δ
torchvision/transforms/functional_tensor.py	64.61% <90.47%> (-0.48%)	⬇️
torchvision/transforms/transforms.py	77.48% <0.00%> (ø)
torchvision/io/image.py	48.38% <0.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6fe11d5...f7448d7. Read the comment docs.

[fmassa]

Implementation looks great, thanks a lot @vfdev-5 !

[fmassa]

Can you also add tests for negative padding?

[fmassa]

Actually, this doesn't work as expected if padding is negative (in which case we crop the image).

[vfdev-5]

Added check and raise an error if padding is negative for symmetric mode

[vfdev-5]

Benchmarking current implementation vs torch.nn.functional.pad (float32 data on CUDA), vs torchvision.transforms.functional.pad (uint8 data on CPU):

Benchmark 1)

Hardware: NVidia 1080Ti

%%timeit -r 7 -n 500
torch.nn.functional.pad(img, padding, mode="reflect").sum().item()
> 102 µs ± 3.66 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)

%%timeit -r 7 -n 500
pad_symmetric(x, padding).sum().item()
344 µs ± 3.85 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)

almost half time is spent in index creation.

Benchmark 2) CPU

%%timeit -r 7 -n 500
pad_symmetric(x_cpu, padding)
> 1.32 ms ± 106 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)

%%timeit -r 7 -n 500
pad(img_pil, padding, padding_mode="symmetric")
726 µs ± 30.2 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)

from typing import List, Tuple

import torch
from torch import Tensor

def pad_symmetric(img: Tensor, padding: List[int]) -> Tensor:
    # padding is left, right, top, bottom
    in_sizes = img.size()

    x_indices = [i for i in range(in_sizes[-1])]  # [0, 1, 2, 3, ...]
    left_indices = [i for i in range(padding[0] - 1, -1, -1)]  # e.g. [3, 2, 1, 0]
    right_indices = [-(i + 1) for i in range(padding[1])]  # e.g. [-1, -2, -3]
    x_indices = torch.tensor(left_indices + x_indices + right_indices)

    y_indices = [i for i in range(in_sizes[-2])]
    top_indices = [i for i in range(padding[2] - 1, -1, -1)]
    bottom_indices = [-(i + 1) for i in range(padding[3])]
    y_indices = torch.tensor(top_indices + y_indices + bottom_indices)

    ndim = img.ndim
    if ndim == 3:
        return img[:, y_indices[:, None], x_indices[None, :]]
    elif ndim == 4:
        return img[:, :, y_indices[:, None], x_indices[None, :]]
    else:
        raise RuntimeError("Symmetric padding of N-D tensors are not supported yet")

device = "cuda"
x = torch.rand(3, 512, 512).to(device)

img = x.unsqueeze(0)

As discussed privately with @fmassa, an optimization is possible (copying the idea from numpy/pytorch):

out = torch.empty(out_shape)
out[..., i0:i1, j0:j2] = input
out[..., border_y, border_x] = input[..., y, x]

This optimization may be done later...

[fmassa]

Thanks a lot!