Remove remaining __torch_function__ calls from regular transform pipeline

[pmeier]

On our journey to micro-optimize transforms v2, we recently merged #6681. This significantly reduced the number of __torch_function__ calls inside the augmentation pipeline. However, we still have some unnecessary calls to it that originate from this idiom that we have on the feature methods:

vision/torchvision/prototype/features/_image.py

Lines 125 to 126 in 149edda

	def horizontal_flip(self) -> Image:
	output = self._F.horizontal_flip_image_tensor(self)

The kernel, here horizontal_flip_image_tensor, is called with self, i.e. a features.Image. Internally, the first operation will thus go through __torch_function__ and unwrap it. Meaning, we are going through the whole protocol for no reason, since the kernel per definition works with plain tensors.

We could simply unwrap, i.e. self.as_subclass(torch.Tensor) before we call the kernel. That would make our implementation a little more verbose, but would prevent any __torch_function__ calls when using our transformations whatsoever. Of course we could also have a self.unwrap() method to reduce the boilerplate. Another option is to move the kernel call under the DisableTorchFunction context manager, similar to what we did in #6681.

from time import perf_counter_ns

import torch
from torch._C import DisableTorchFunction
from torchvision.prototype import features
from torchvision.prototype.transforms import functional as F


tensor = torch.rand(3, 256, 256)
image = features.Image(tensor)

arbitrary_unary_image_tensor_kernel = F.autocontrast_image_tensor


def minimum():
    arbitrary_unary_image_tensor_kernel(tensor)


def unwrap_boilerplate():
    arbitrary_unary_image_tensor_kernel(image.as_subclass(torch.Tensor))


def _unwrap_helper(feature):
    return feature.as_subclass(torch.Tensor)


def unwrap_with_helper_method():
    arbitrary_unary_image_tensor_kernel(_unwrap_helper(image))


def context_manager():
    with DisableTorchFunction():
        arbitrary_unary_image_tensor_kernel(image)


def baseline():
    arbitrary_unary_image_tensor_kernel(image)


for fn in [minimum, unwrap_boilerplate, unwrap_with_helper_method, context_manager, baseline]:
    for _ in range(1_000):
        fn()

    time_diffs_ns = []
    for _ in range(100_000):
        start = perf_counter_ns()
        fn()
        stop = perf_counter_ns()
        time_diffs_ns.append(stop - start)

    time_diffs = torch.tensor(time_diffs_ns, dtype=torch.float64) * 1e-9
    print(f"{fn.__name__}: median={time_diffs.median() * 1e6:.2f} µs, std={time_diffs.std() * 1e6:.2f} µs")

minimum: median=69.79 µs, std=3.61 µs
unwrap_boilerplate: median=74.35 µs, std=5.30 µs
unwrap_with_helper_method: median=74.65 µs, std=5.57 µs
context_manager: median=76.62 µs, std=3.58 µs
baseline: median=94.84 µs, std=4.42 µs

Benchmarking shows that unwrapping is slightly faster, but the difference to the context manager is small enough to also opt for consistency with #6681 and use it instead.

cc @vfdev-5 @datumbox @bjuncek

[pmeier]

On second thought, I would prefer the unwrap method over using the context manager for three reasons:

1. [proto] Reduce number of calls of __torch_function__ #6681 (comment) proposed to add functionality to the __torch_function__ protocol to ignore certain functions. If that is implemented in core, we no longer need the properties for the attributes since we could simply ignore __torch_function__ for them. In turn the context manager would also be removed. Thus, using the context manager in our other methods would align them with something that might be removed in the future.
2. I think the idiom return features.Image.wrap_like(self, some_kernel(self.unwrap())) is actually quite nice to hammer home the point that our kernels don't need our features.
3. The combination of wrap_like and unwrap is easier to understand than wrap_like and DisableTorchFunction

[datumbox]

IMO adding a public unwrap() method to just do self.as_subclass(torch.Tensor) seems unnecessary. It also exposes publicly this implementation detail. I would be in favour of a PR that does the above call before invoking the Kernel to eliminate any unnecessary __torch_function__ calls. If the Core team improves the speed on the future, we could remove the verbosity on our side.

[pmeier]

IMO adding a public unwrap() method to just do self.as_subclass(torch.Tensor) seems unnecessary. It also exposes publicly this implementation detail.

IMO having this public is a good thing, because it is likely not obvious for users if they have a feature how to get the underlying plain tensor. My guess is that they will just use image.data, but that is soft-deprecated.

I would be in favour of a PR that does the above call before invoking the Kernel to eliminate any unnecessary __torch_function__ calls.

Meaning you are in favor of self.as_subclass(torch.Tensor), correct?

@vfdev-5 given that you have reacted negative to #6781 (comment) can you get behind Vasilis favorite or do you want the context manager? I hope it is the former or otherwise the three of us are advocating for something different 😄

[vfdev-5]

I have an impression that this proposal looks like "remove tensor subclassing". I also agree that unwrap method is exposing implementation details. Finally, I think it is unnecessary to add new public API method to gain ~10us.

because it is likely not obvious for users if they have a feature how to get the underlying plain tensor.

Docs should explain that all features are subclasing tensors -> behave like tensors and there is no particular need to get underlying tensor.

[datumbox]

IMO having this public is a good thing, because it is likely not obvious for users if they have a feature how to get the underlying plain tensor.

For every new public method we introduce on the subclassed tensors, we run the risk on the future to get name collisions with core methods. Core already offers a way to "unwrap" a tensor and that's the as_subclass(torch.Tensor) idiom. We should use that.

The optimizations we do right now, should be implementation details to circumvent speed issues. As Core improves the work on the Tensor Subclassing, these might become unnecessary on the future and we might be able to remove them. Thus I am in favour of going for the slightly more verbose idiom to cut down the overhead of the new API but keep this private.