Performance improvements for transforms v2 vs. v1

[pmeier]

In addition to a lot of other goodies that transforms v2 will bring, we are also actively working on improving the performance. This is a tracker / overview issue of our progress.

Performance was measured with this benchmark script. Unless noted otherwise, the performance improvements reported above were computed on uint8, RGB images and videos while running single-threaded on CPU. You can find the full benchmark results alongside the benchmark script. The results will be constantly updated if new PRs are merged that have an effect on the kernels.

Kernels:

• color
  • adjust_brightness [proto] Speed up adjust color ops #6784
  • adjust_contrast [proto] Speed up adjust color ops #6784 [prototype] Speed up adjust_contrast_image_tensor #6933
  • adjust_gamma [prototype] Speed improvement for adjust gamma op #6820 replace tensor division with scalar division and tensor multiplication #6903
  • adjust_hue [proto] Speed improvements for adjust hue op #6805 replace tensor division with scalar division and tensor multiplication #6903 [prototype] Speed up adjust_hue_image_tensor #6938
  • adjust_saturation [proto] Speed up adjust color ops #6784 [prototype] Minor change on adjust_saturation_image_tensor uint8 #6940
  • adjust_sharpness [proto] Speed up adjust color ops #6784 [prototype] Speed up adjust_sharpness_image_tensor #6930
  • autocontrast [proto] Speed improvement for autocontrast op #6811 [prototype] Speed up autocontrast_image_tensor #6935 [prototype] Port elastic and minor cleanups #6942
  • equalize [proto] Small improvement for tensor equalize op #6738, [proto] Performance improvements for equalize op #6757, another round of perf improvements for equalize #6776
  • invert improve performance of {invert, solarize}_image_tensor #6819
  • posterize remove unneccesary checks from posterize_image_tensor #6823, extend support of posterize to all integer and floating dtypes #6847
  • solarize improve performance of {invert, solarize}_image_tensor #6819
• geometry
  • affine Fix bug on prototype pad #6949
  • center_crop [prototype] Optimize Center Crop performance #6880 Fix bug on prototype pad #6949
  • crop Fix bug on prototype pad #6949
  • elastic [prototype] Port elastic and minor cleanups #6942
  • erase [prototype] Remove _FT aliases from functional #6983
  • five_crop: Composite kernel Fix bug on prototype pad #6949
  • pad Fix bug on prototype pad #6949
  • perspective [proto] Small optim for perspective op on images #6907 Fix bug on prototype pad #6949
  • resize [prototype] Clean up and port the resize kernel in V2 #6892
  • resized_crop: Composite kernel [prototype] Clean up and port the resize kernel in V2 #6892 Fix bug on prototype pad #6949
  • rotate Fix bug on prototype pad #6949
  • ten_crop: Composite kernel Fix bug on prototype pad #6949
• meta
  • convert_color_space [proto] Speed up adjust color ops #6784 [prototype] Minor improvements on functional #6832
  • convert_dtype improve perf on convert_image_dtype and add tests #6795 replace tensor division with scalar division and tensor multiplication #6903
    • There is still some performance gain left for int to int conversion. Currently, we are using a multiplication
      but theoretically bit shifts are faster. However, on PyTorch core the CPU kernels for bit shifts are not
      vectorized making them slower for regular sized images than a multiplication. Vectorized CPU code implementing left shift operator. pytorch#88607
• misc
  • gaussian_blur [proto] Small optimization for gaussian_blur functional op #6762 [prototype] Gaussian Blur clean up #6888
  • normalize [prototype] Speed improvement for normalize op #6821

Transform Classes:

• MixUp/CutMix [prototype] Speed up Augment Transform Classes #6835
• ColorJitter, RandomPhotometricDistort [prototype] Minor speed and nit optimizations on Transform Classes #6837

C++ (PyTorch core):

• vertical_flip [prototype] Remove _FT aliases from functional #6983 Optimized vertical flip using memcpy pytorch#89414
• horizontal_flip [prototype] Remove _FT aliases from functional #6983 Vectorized horizontal flip implementation pytorch#88989 Optimized vertical flip using memcpy pytorch#89414

cc @vfdev-5 @datumbox @bjuncek

[datumbox]

Concerning elastic and all the affine transform kernels (affine, perspective, rotate), there are some very limited opportunities for optimization. Perhaps a couple of in-place ops in elastic_transform & _perspective_grid and a few optimizations in _apply_grid_transform (split of mask and img, bilinear fill estimation etc). Also some minor fixes related to the input assertion. @vfdev-5 would you be OK to assess on your side whether it make sense to do these or leave the methods on _FT to avoid copy-pasting? Perhaps you have in mind other optimizations that I can't see that could affect performance?

Concerning crop, erase, pad, resize, horizontal_flip and vertical_flip, I don't see any further improvements other than the input assertions. It might be worth to have a look on your side, @pmeier and @vfdev-5, incase you see something I don't.

[pmeier]

I did another deep dive into the ops in the second paragraph of #6818 (comment) and I'm fairly confident that there is little we can do to improve on our side. The only two things I found are

• For padding modes "edge" and "reflect" we cast to float32 and back:

  vision/torchvision/transforms/functional_tensor.py

  Lines 426 to 431 in c84dbfa

  	if (padding_mode != "constant") and img.dtype not in (torch.float32, torch.float64):
  	# Here we temporary cast input tensor to float
  	# until pytorch issue is resolved :
  	# https://github.com/pytorch/pytorch/issues/40763
  	need_cast = True
  	img = img.to(torch.float32)

  There is a long standing issue on PyTorch core Enhance supported types of functional.pad  pytorch#40763 that reports this and is assigned to @vfdev-5.

• We support "symmetric" padding in F.pad, but torch.nn.functional.pad doesn't. Thus, we have a custom implementation for it
  

  vision/torchvision/transforms/functional_tensor.py

  Line 330 in c84dbfa

  def _pad_symmetric(img: Tensor, padding: List[int]) -> Tensor:

  
  Since it is written in Python, a possible speed up would be to implement this padding mode in C++ on the PyTorch core side.

Fixing this, we would get speed-ups for padding modes "edge", "reflect", and "symmetric" but not for the default and ubiquitous "constant" padding mode. Skimming the repository, it seems the only time we use non-"constant" padding is

vision/torchvision/transforms/functional_tensor.py

Line 762 in c84dbfa

img = torch_pad(img, padding, mode="reflect")

In there the image is guaranteed to be float and thus would not get any performance boost.

While I think both things mentioned above would be good to have in general, I don't think we should prioritize them.

[vfdev-5]

Concerning elastic and all the affine transform kernels (affine, perspective, rotate), there are some very limited opportunities for optimization. Perhaps a couple of in-place ops in elastic_transform & _perspective_grid and a few optimizations in _apply_grid_transform (split of mask and img, bilinear fill estimation etc). Also some minor fixes related to the input assertion. @vfdev-5 would you be OK to assess on your side whether it make sense to do these or leave the methods on _FT to avoid copy-pasting? Perhaps you have in mind other optimizations that I can't see that could affect performance?

Checking various options with affine, there is no obvious way to improve runtime perfs. Yes, we can make some inplace "split of mask and img, bilinear fill estimation etc". There is also an open issue about incorrect behaviour of bilinear mode with provided not-None fill (#6517). Given that I think we can keep this implementation.

[vadimkantorov]

About not vectorized bitwise shifts, is there an issue in core?

[pmeier]

About not vectorized bitwise shifts, is there an issue in core?

I don't think so, but @alexsamardzic wanted to have a look at it.

Edit: pytorch/pytorch#88607