Repository: pytorch
Issue: bilinear interpolate is very slow under mixed precision training mode.
Labels: module: performance,module: nn,module: cuda,triaged,module: amp (automated mixed precision),module: interpolation

## 🐛 Bug

When training models with mixed precision mode, interpolate operation is very slow.

## To Reproduce


Please test this piece of code:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.cuda.amp as amp
import torchvision

import numpy as np
import random
import time
import datetime

torch.manual_seed(15)
random.seed(15)
np.random.seed(15)
torch.backends.cudnn.deterministic = True


inter = 'bilinear'

class Model(nn.Module):
    def __init__(self, n_classes):
        super(Model, self).__init__()
        net = torchvision.models.resnet18(pretrained=False)
        self.conv1 = net.conv1
        self.bn1 = net.bn1
        self.maxpool = net.maxpool
        self.relu = net.relu
        self.layer1 = net.layer1
        self.layer2 = net.layer2
        self.layer3 = net.layer3
        self.layer4 = net.layer4
        up_factor = 32
        if inter == 'pixshuffle':
            n_classes = up_factor * up_factor * n_classes
            self.fc = nn.Sequential(
                nn.Conv2d(512, n_classes, 1, 1, 0),
                nn.PixelShuffle(up_factor)
                    )
        elif inter == 'bilinear':
            self.fc = nn.Sequential(
                nn.Conv2d(512, n_classes, 1, 1, 0),
                nn.Upsample(scale_factor=up_factor, mode='bilinear', align_corners=False)
                    )
    def forward(self, x):
        feat = self.conv1(x)
        feat = self.bn1(feat)
        feat = self.relu(feat)
        feat = self.maxpool(feat)
        feat = self.layer1(feat)
        feat = self.layer2(feat)
        feat = self.layer3(feat)
        feat = self.layer4(feat)
        feat = self.fc(feat)
        return feat

c = 19
bs, h, w = 8, 640, 640
net = Model(c)
criteria = nn.CrossEntropyLoss()
net.cuda()
net.train()
criteria.cuda()

optim = torch.optim.SGD(net.parameters(), lr=1e-2)
scaler = amp.GradScaler()

t1 = time.time()
for it in range(100):
    inten = torch.randn(bs, 3, h, w).cuda()#.half()
    lbs = torch.randint(0, c, (bs, h, w)).cuda()

    optim.zero_grad()
    with amp.autocast(enabled=True):
        logits = net(inten)
        loss = criteria(logits, lbs)
    scaler.scale(loss).backward()
    scaler.step(optim)
    scaler.update()

t2 = time.time()
t_str = str(datetime.timedelta(seconds=t2 - t1))
print('time of program is: ', t_str)
```
By assigning inter='pixshuffle', the program runs for 26s on my platform, but by assigning inter='bilinear', the program runs for 54s.

## Expected behavior

The time gap should be not so much 

## Environment

```
Collecting environment information...
PyTorch version: 1.8.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.17.1
Libc version: glibc-2.27

Python version: 3.8.8 | packaged by conda-forge | (default, Feb 20 2021, 16:22:27)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-3.10.0-862.el7.x86_64-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: Tesla T4
GPU 1: Tesla T4
GPU 2: Tesla T4
GPU 3: Tesla T4
GPU 4: Tesla T4
GPU 5: Tesla T4
GPU 6: Tesla T4
GPU 7: Tesla T4

Nvidia driver version: 450.80.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-loss==0.0.0
[pip3] torch==1.8.1
[pip3] torchvision==0.9.0
[conda] blas                      1.0                         mkl    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] cudatoolkit               10.2.89              h8f6ccaa_8    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
[conda] mkl                       2020.2                      256    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] mkl-service               2.3.0            py38h1e0a361_2    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
[conda] mkl_fft                   1.3.0            py38h54f3939_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] mkl_random                1.2.0            py38hc5bc63f_1    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
[conda] numpy                     1.19.2           py38h54aff64_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] numpy-base                1.19.2           py38hfa32c7d_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] pytorch                   1.8.1           py3.8_cuda10.2_cudnn7.6.5_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
[conda] pytorch-loss              0.0.0                    pypi_0    pypi
[conda] torchvision               0.9.0                    pypi_0    pypi

```


## Additional context

I also tried nearest interpolation which is as fast as pixel-shuffle. Do not understand why bilinear can be slower in such extend.


cc @VitalyFedyunin @ngimel @heitorschueroff @albanD @mruberry @jbschlosser @mcarilli @ptrblck