Hack to improve performance of resize op with nearest mode on 2D

[vfdev-5]

Benchmarks:

Main (Feature vs PIL)

[--- Resize measurements: 500x600 -> 256 ---]
                         |  stable  |    v2  
1 threads: ----------------------------------
      PIL mask data      |   55.4   |        
      Feature Mask data  |          |  1009.3

Times are in microseconds (us).

[--- Resize measurements: 500x600 -> 520 ---]
                         |  stable  |    v2  
1 threads: ----------------------------------
      PIL mask data      |  178.0   |        
      Feature Mask data  |          |  3773.3

Times are in microseconds (us).

[--- Resize measurements: 500x600 -> 720 ---]
                         |  stable  |    v2  
1 threads: ----------------------------------
      PIL mask data      |  320.3   |        
      Feature Mask data  |          |  7101.1

Times are in microseconds (us).

This PR (Feature vs PIL)

[-- Resize measurements: 500x600 -> 256 ---]
                         |  stable  |    v2
1 threads: ---------------------------------
      PIL mask data      |   53.4   |
      Feature Mask data  |          |  293.0

Times are in microseconds (us).

[-- Resize measurements: 500x600 -> 520 ---]
                         |  stable  |    v2
1 threads: ---------------------------------
      PIL mask data      |  165.9   |
      Feature Mask data  |          |  652.5

Times are in microseconds (us).

[--- Resize measurements: 500x600 -> 720 ---]
                         |  stable  |    v2
1 threads: ----------------------------------
      PIL mask data      |  297.3   |
      Feature Mask data  |          |  1105.2

Times are in microseconds (us).

Code:

import numpy as np
import PIL
from functools import partial

import torch
import torch.utils.benchmark as benchmark

from torchvision.transforms import functional as F_stable
from torchvision.transforms.functional import InterpolationMode
from torchvision.prototype import features
from torchvision.prototype.transforms import functional as F_v2


def get_pil_mask(size):
    target_data = np.zeros(size, dtype="int32")
    target_data[110:140, 120:160] = 1
    target_data[10:40, 120:160] = 2
    target_data[110:140, 20:60] = 3
    target_data[size[0] // 2 : size[0] // 2 + 50, size[1] // 2 : size[1] // 2 + 60] = 4
    target = PIL.Image.fromarray(target_data).convert("L")
    return target


def main():
    results = []
    min_run_time = 2

    for size in [256, 520, 720]:
        pil_mask = get_pil_mask((500, 600))
        mask = features.Mask(F_v2.pil_to_tensor(pil_mask).squeeze(0))

        transform_stable = partial(F_stable.resize, size=size, interpolation=InterpolationMode.NEAREST)
        transform_v2 = partial(F_v2.resize, size=[size], interpolation=InterpolationMode.NEAREST)

        # PIL resize
        results.append(
            benchmark.Timer(
                stmt=f"transform(data)",
                globals={
                    "data": pil_mask,
                    "transform": transform_stable,
                },
                num_threads=torch.get_num_threads(),
                label=f"Resize measurements: 500x600 -> {size}",
                sub_label="PIL mask data",
                description="stable",
            ).blocked_autorange(min_run_time=min_run_time)
        )
        # Mask resize
        results.append(
            benchmark.Timer(
                stmt=f"transform(data)",
                globals={
                    "data": mask,
                    "transform": transform_v2,
                },
                num_threads=torch.get_num_threads(),
                label=f"Resize measurements: 500x600 -> {size}",
                sub_label="Feature Mask data",
                description="v2",
            ).blocked_autorange(min_run_time=min_run_time)
        )

    compare = benchmark.Compare(results)
    compare.print()


if __name__ == "__main__":
    main()

Using this hack we can reduce slowdown on segmentation dataaug pipeline (PIL seg mask vs features.Mask) from x10 to x2 :

## https://github.com/vfdev-5/tvapiv2_benchmarks
# python -u main.py segmentation --with_time --single_dtype=PIL
Timestamp: 20220929-110049
Torch version: 1.13.0.dev20220906+cu113
Torchvision version: 0.14.0a0
Num threads: 1

[ Segmentation transforms measurements ]
                      |  stable  |   v2
1 threads: -----------------------------
      PIL Image data  |   6.3    |  11.8

Times are in milliseconds (ms).


-----

-- Benchmark: Segmentation
- Stable transforms: RefCompose(
    <seg_transforms.RandomResize object at 0x7f0ad007f1c0>
    <seg_transforms.RandomHorizontalFlip object at 0x7f0ad007f340>
    <seg_transforms.RandomCrop object at 0x7f0ad007f3a0>
    <seg_transforms.PILToTensor object at 0x7f0ad007f4f0>
    <seg_transforms.ConvertImageDtype object at 0x7f0ad006d2b0>
    <seg_transforms.Normalize object at 0x7f0ad006d9a0>
)
- Transforms v2: Compose(
      SegWrapIntoFeatures()
      RandomResize(min_size=260, max_size=1040, interpolation=InterpolationMode.BILINEAR)
      RandomHorizontalFlip(p=1.0)
      PadIfSmaller(size=480)
      RandomCrop(size=(480, 480), pad_if_needed=False, padding_mode=constant)
      ToImageTensor()
      ConvertImageDtype()
      Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], inplace=False)
)

[datumbox]

Good find! We should eventually fix on PyTorch Core but it will do for now.

Shall we move the method directly in prototype? This will allow us to simplify some of the code and it will permit the easier comparisons of the performance improvements made on Transforms V2. We are also extremely close to cutting the branch of the release, so we should avoid merging hacks on main. WDYT?

[vfdev-5]

Sounds good to put things on prototype.
However, JIT is not passing, I'll be searching for a workaround.

[datumbox]

The very last resort for JIT is to make it available only when not scripting.

[NicolasHug]

Unless we're in a rush, wouldn't it be better to wait for a proper solution in core? Something similar to pytorch/pytorch#83840 (comment)

I agree that with the approaching release we may not want to risk too much

[datumbox]

We need it sooner than later but to avoid issues we will move it to prototype. It is likely there will be many workarounds like this to close the gaps between V1 and V2 in terms of speed and make them usable for the TorchMultimodal team. Most of them will be removed long before the API moved out of prototype.

[NicolasHug]

FWIW, moving the optimization into the prototype area won't really give us a fair comparison between V1 and V2, since the optimization can be applied to both versions.

[vfdev-5]

while moving the hack to prototype we have a slow down due to the fact that we cast non-float data to float (and copy), this makes input contiguous and we have to process twice more data as we expanded (N, 1->2, H, W).

[datumbox]

What if the input is (N, 1, H, W), do you still need to do this hack? Will your code work in this scenario? What's the performance gain for batch of say 32?

[vfdev-5]

Yes, you are right, I need to use N and not 1

[vfdev-5]

It gives ~4x speedup if using batch=32:

Timestamp: 20220928-181325
Torch version: 1.13.0.dev20220906+cu113
Torchvision version: 0.14.0a0
Num threads: 1

[--------------------------------- Mask Resize measurements --------------------------------]
                                       |  Original (slow) mask 2d  |  Hacked (faster) mask 2d
1 threads: ----------------------------------------------------------------------------------
      ([32, 1, 500, 500]), 500 -> 128  |            5.5            |            1.6          

Times are in milliseconds (ms).

[datumbox]

LGTM, thanks!