pytorch istft runs slower than torchaudio istft especially at higher n_fft

🐛 Bug

There seems to be a significant performance drop after migrating the istft from torchaudio to pytorch.
migration discussed here #34827
I’ve also noticed the performance gap increases with greater n_fft.

To Reproduce

pip install torch==1.6.0

import torch
import time
n_fft = 2**15
x = torch.rand(size=[1, n_fft//2 + 1, 10, 2])
start = time.time()
y = torch.functional.istft(x,n_fft)
end = time.time()
print(end - start)

pip install torchaudio==0.5.1

import torch
import torchaudio
import time
n_fft = 2**15
x = torch.rand(size=[1, n_fft//2 + 1, 10, 2])
start = time.time()
y = torchaudio.functional.istft(x,n_fft)
end = time.time()
print(end - start)

Expected behavior

I would expect similar timings between:

torch.functional.istft in torch==1.6.0
and
torchaudio.functional.istft in torchaudio==0.5.1

Environment running torch.functional.istft

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] torch==1.6.0
[pip3] torchaudio==0.6.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.6.1+cu101
[conda] Could not collect

Environment running torchaudio.functional.istft

PyTorch version: 1.5.1+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] torch==1.5.1+cu101
[pip3] torchaudio==0.5.1
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.6.1+cu101
[conda] Could not collect

Additional context

Both environments are from google colab

cc @ezyang @gchanan @zou3519 @VitalyFedyunin @ngimel

1 possible answer(s) on “pytorch istft runs slower than torchaudio istft especially at higher n_fft

  1. I confirm that I could reproduce the issue.

    Total Par Call
    PyTorch 1.6.0 9.0404 0.904
    torchaudio 0.5.0 0.1052 0.010
    script
    import time
    
    import torch
    import torchaudio
    
    n_fft = 2**15
    x = torch.rand(size=[1, n_fft//2 + 1, 10, 2])
    
    start = time.monotonic()
    n = 10
    for _ in range(n):
        # y = torchaudio.functional.istft(x,n_fft)
        y = torch.functional.istft(x,n_fft)
    elapsed = time.monotonic() - start
    
    print(elapsed, elapsed / n)
    Env 1: PyTorch 1.6.0
    $ python -m torch.utils.collect_env
    Collecting environment information...
    PyTorch version: 1.5.0
    Is debug build: No
    CUDA used to build PyTorch: 10.2
    
    OS: Ubuntu 18.04.3 LTS
    GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
    CMake version: version 3.10.2
    
    Python version: 3.6
    Is CUDA available: No
    CUDA runtime version: Could not collect
    GPU models and configuration:
    GPU 0: Quadro GP100
    GPU 1: Quadro GP100
    
    Nvidia driver version: 418.116.00
    cuDNN version: Could not collect
    
    Versions of relevant libraries:
    [pip3] numpy==1.18.5
    [pip3] torch==1.5.0
    [pip3] torchaudio==0.5.0a0+3305d5c
    [conda] blas                      1.0                         mkl
    [conda] mkl                       2020.1                      217
    [conda] mkl-service               2.3.0            py36he904b0f_0
    [conda] mkl_fft                   1.1.0            py36h23d657b_0
    [conda] mkl_random                1.1.1            py36h0573a6f_0
    [conda] pytorch                   1.5.0           py3.6_cuda10.2.89_cudnn7.6.5_0    pytorch
    [conda] torchaudio                0.5.0                      py36    pytorch
    
    Env 2: torchaudio 0.5.0
    $ python -m torch.utils.collect_env
    Collecting environment information...
    PyTorch version: 1.6.0
    Is debug build: No
    CUDA used to build PyTorch: 10.2
    
    OS: Ubuntu 18.04.3 LTS
    GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
    CMake version: version 3.10.2
    
    Python version: 3.6
    Is CUDA available: No
    CUDA runtime version: Could not collect
    GPU models and configuration:
    GPU 0: Quadro GP100
    GPU 1: Quadro GP100
    
    Nvidia driver version: 418.116.00
    cuDNN version: Could not collect
    
    Versions of relevant libraries:
    [pip3] numpy==1.18.5
    [pip3] torch==1.6.0
    [pip3] torchaudio==0.6.0a0+f17ae39
    [conda] blas                      1.0                         mkl
    [conda] cudatoolkit               10.2.89              hfd86e86_1
    [conda] mkl                       2020.1                      217
    [conda] mkl-service               2.3.0            py36he904b0f_0
    [conda] mkl_fft                   1.1.0            py36h23d657b_0
    [conda] mkl_random                1.1.1            py36h0573a6f_0
    [conda] numpy                     1.18.5           py36ha1c710e_0
    [conda] numpy-base                1.18.5           py36hde5b4d6_0
    [conda] pytorch                   1.6.0           py3.6_cuda10.2.89_cudnn7.6.5_0    pytorch
    [conda] torchaudio                0.6.0                      py36    pytorch