Repository: pytorch
Issue: tensor subclass and `__torch_function__` performance issues
Labels: triaged,module: __torch_function__

### 🐛 Describe the bug

Sorry If i'm submitting this as the wrong issue type. I would like to subclass `torch.Tensor` to add some additional data/functionality for my use case. This is useful because I would rather have the data encapsulated inside the object rather than having to pass around tuples of tuples of... For example this can be very useful to store meta data, with the advantage of not needing to change existing interfaces in order to accommodate new data being added to the subclass class.

I decided to run some experiments as I was seeing some performance issues. Here is the code to reproduce my experiments. I opted for a bare minimum tensor subclass just to rule out that my code was causing the performance hit. However the actual use case I would like to do is to modify the `__torch_function__` to process `extra_data` with different operations.

```python
from time import sleep, time
import torch
from torchvision.models.mobilenetv3 import mobilenet_v3_large, mobilenet_v3_small
from torchvision.models.efficientnet import efficientnet_b2
from torchvision.models.vgg import vgg16, vgg11
from torchvision.models.alexnet import alexnet
from torchvision.models.googlenet import googlenet

class MyTensor(torch.Tensor):

    extra_data: str = ""

    def __repr__(self):
        return f"extra_data: {self.extra_data}\ndata:\n{super().__repr__()}"

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        return super().__torch_function__(func, types, args, kwargs)


def do_opt(model, img: torch.Tensor, opt: torch.optim.Optimizer):
    result = model(img)

    opt.zero_grad()
    # just use any loss function
    loss = torch.mean(torch.square(result - 1))
    loss.backward()
    opt.step()

def run_test(model_func, n=100, burn_in=5):
    print(f"testing model {model_func.__name__}")

    model = model_func().cuda()

    img = torch.rand([32, 3, 256, 256]).cuda()

    opt = torch.optim.Adam(model.parameters())

    for i in range(burn_in):
        do_opt(model, img, opt)

    start1 = time()
    for i in range(n):
        do_opt(model, img, opt)
    end1 = time()

    img = img.as_subclass(MyTensor)
    img.extra_data = "test"

    start2 = time()
    for i in range(n):
        do_opt(model, img, opt)
    end2 = time()

    tensor_time = 1000 * (end1 - start1) / n
    mytensor_tme = 1000 * (end2 - start2) / n
    pct = 100 * (mytensor_tme - tensor_time) / tensor_time
    sign = "+" if pct >= 0 else ""

    print(f"torch.Tensor: {tensor_time:0.2f}ms per iteration")
    print(f"    MyTensor: {mytensor_tme:0.2f}ms per iteration")
    print(f"% difference: {sign}{pct:0.2f}%")

def main():

    models_to_test = [
        mobilenet_v3_large,
        mobilenet_v3_small,
        efficientnet_b2,
        vgg16,
        vgg11,
        alexnet
    ]

    for model_func in models_to_test:
        run_test(model_func)

if __name__ == "__main__":
    main()
```

Here is the output with pytorch nightly.

```
testing model mobilenet_v3_large
torch.Tensor: 54.04ms per iteration
    MyTensor: 66.64ms per iteration
% difference: +23.31%
testing model mobilenet_v3_small
torch.Tensor: 25.54ms per iteration
    MyTensor: 31.23ms per iteration
% difference: +22.25%
testing model efficientnet_b2
torch.Tensor: 144.74ms per iteration
    MyTensor: 169.72ms per iteration
% difference: +17.26%
testing model vgg16
torch.Tensor: 286.61ms per iteration
    MyTensor: 324.19ms per iteration
% difference: +13.11%
testing model vgg11
torch.Tensor: 161.03ms per iteration
    MyTensor: 181.56ms per iteration
% difference: +12.75%
testing model alexnet
torch.Tensor: 26.12ms per iteration
    MyTensor: 27.51ms per iteration
% difference: +5.29%
```
Here is the output with pytorch stable:
```
testing model mobilenet_v3_large
torch.Tensor: 55.69ms per iteration
    MyTensor: 67.54ms per iteration
% difference: +21.27%
testing model mobilenet_v3_small
torch.Tensor: 25.23ms per iteration
    MyTensor: 29.33ms per iteration
% difference: +16.26%
testing model efficientnet_b2
torch.Tensor: 142.44ms per iteration
    MyTensor: 169.52ms per iteration
% difference: +19.01%
testing model vgg16
torch.Tensor: 291.71ms per iteration
    MyTensor: 330.15ms per iteration
% difference: +13.18%
testing model vgg11
torch.Tensor: 163.73ms per iteration
    MyTensor: 185.06ms per iteration
% difference: +13.03%
testing model alexnet
torch.Tensor: 26.09ms per iteration
    MyTensor: 27.48ms per iteration
% difference: +5.33%
```

As you can see in some cases the the subclassed tensor can add an overhead, with around 20% more time. I haven't tested this in a real world situation yet, but when I do I'll add more to this issues. Hopefully it won't make training take 20% more time, which could be significant if it takes 24 hours to train a model.

### Versions

I ran these tests with two versions of pytorch.

Nightly:

```
Collecting environment information...
PyTorch version: 1.12.0.dev20220307+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-6ubuntu2) 7.5.0
Clang version: Could not collect
CMake version: version 3.18.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-100-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.0.194
GPU models and configuration: GPU 0: GeForce RTX 2080 Ti
Nvidia driver version: 460.73.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.1
[pip3] mypy==0.800
[pip3] mypy-extensions==0.4.3
[pip3] mypy-protobuf==1.23
[pip3] numpy==1.21.3
[pip3] pytorch-lightning==1.4.4
[pip3] torch==1.12.0.dev20220307+cu113
[pip3] torchaudio==0.12.0.dev20220307+cu113
[pip3] torchmetrics==0.5.0
[pip3] torchvision==0.13.0.dev20220307+cu113
[conda] Could not collect
```

and stable:
```
Collecting environment information...
PyTorch version: 1.10.2+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-6ubuntu2) 7.5.0
Clang version: Could not collect
CMake version: version 3.18.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-100-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.0.194
GPU models and configuration: GPU 0: GeForce RTX 2080 Ti
Nvidia driver version: 460.73.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.4
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.1
[pip3] mypy==0.800
[pip3] mypy-extensions==0.4.3
[pip3] mypy-protobuf==1.23
[pip3] numpy==1.21.3
[pip3] pytorch-lightning==1.4.4
[pip3] torch==1.10.2+cu113
[pip3] torchaudio==0.10.2+cu113
[pip3] torchmetrics==0.5.0
[pip3] torchvision==0.11.3+cu113
[conda] Could not collect
```

cc @hameerabbasi @rgommers @peterbell10