GPU computations weirdly slowed down if calling stft

JuanFMontesinos: <p>Hi,<br>
I was training some audiovisual source separation models.<br>
I was computing STFT in the GPU for speeding up the preprocessing. The following code is a wrapper around a nn.Module to apply some preprocessing.</p>
<p>The relevant code is:</p>
<pre><code class="lang-auto">    def forward(self, inputs: dict):
        """
        Inputs contains the following keys:
           mixture: the mixture of audios as spectrogram
           sk: skeletons of shape N,C,T,V,M which is batch, channels, temporal, joints, n_people
           video: video
        """

        def numpy(x):
            return x[0].detach().cpu().numpy()

        self.n_sources = 2
        with torch.no_grad():
            llcp_embeddings = inputs['llcp_embedding'].transpose(1, 2)


            srcm = inputs['audio']
            srcs = inputs['audio_acmt']

            # Computing STFT
            spm = self.wav2sp(srcm)  # Spectrogram main BxFxTx2
            sps = self.wav2sp(srcs)  # Spectrogram secondary BxFxTx2


            sources = [spm, sps]
            sp_mix_raw = sum(sources) / self.n_sources
            # Downsampling to save memory
            spm = spm[:, ::2, ...]
            sps = sps[:, ::2, ...]
            sp_mix = sp_mix_raw[:, ::2, ...]  # BxFxTx2



            x = sp_mix.permute(0, 3, 1, 2)

        pred = self.core_forward(audio_input=x.detach().requires_grad_(),
                                 visual_input=llcp_embeddings)
        return pred


    def core_forward(self, *, audio_input, visual_input):
        outx = self.llcp(audio_input, visual_input)
        output = {'mask': outx, 'ind_end_feats': None, 'visual_features': None}
        return output
</code></pre>
<p>Where <code>self.wav2sp</code> uses <code>Spectrogram</code> operator from torchaudio and <code>self.core_forward</code> just calls the forward code.</p>
<p>I realised that if I call the code cloning the spectrogram, the speed is boosted.</p>
<pre><code class="lang-auto">        pred = self.core_forward(audio_input=x.clone().requires_grad_(),
                                 visual_input=llcp_embeddings)
</code></pre>
<p>Profiling results without cloning</p>
<pre><code class="lang-auto">
Self CPU time total: 8.264s
CUDA time total: 8.264s


Process finished with exit code 0
</code></pre>
<p>Profiling results withing cloning.</p>
<pre><code class="lang-auto">Self CPU time total: 107.938ms
CUDA time total: 107.809ms


Process finished with exit code 0

</code></pre>
<p>Does anyone know why?<br>
I found this occur for any¿? (i tried several models and always happens) model</p>
<p>Profiled like this:</p>
<pre><code class="lang-auto">if __name__=='__main__':
    USE_W = True
    N = 2
    DEVICE = torch.device('cuda:0')
    # DEVICE=torch.device('cpu')



    if USE_W:
        model = LlcpNet(audio_length=65535, audio_samplerate=16384,
                        n_fft=1022, hop_length=256, n_mel=128, sp_freq_shape=1022 // 2 + 1,

                        video_enabled=False, llcp_enabled=True,
                        skeleton_enabled=False, device=DEVICE).to(DEVICE)
        inputs = {'llcp_embedding': torch.rand(N, 100, 512).to(DEVICE),
                  'audio': torch.rand(N, 65535).to(DEVICE),
                  'audio_acmt': torch.rand(N, 65535).to(DEVICE)}
    else:
        model = Llcp().to(DEVICE)
        inputs = {'input_video': torch.rand(N, 512, 100).to(DEVICE),
                  'input_audio': torch.rand(N, 2, 256, 256).to(DEVICE)}
    # input_audio will be (N,2,256,256)
    # input_video will be of size (N,512,100)


    with profiler.profile(with_stack=True, profile_memory=True, use_cuda=True) as prof:
        output = model(inputs) if USE_W else model(**inputs)
    print(prof.key_averages().table(sort_by='self_cpu_time_total', row_limit=-1))
    print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=-1))

</code></pre>
<p>A standalone script to check this can be found at:</p><aside class="onebox githubgist">
  <header class="source">

      <a href="https://gist.github.com/JuanFMontesinos/90db013f3c736a0a702faae28e73f4ff" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/JuanFMontesinos/90db013f3c736a0a702faae28e73f4ff" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/JuanFMontesinos/90db013f3c736a0a702faae28e73f4ff</a></h4>

  <h5>test_issue.py</h5>
  <pre><code class="Python">from functools import partial

import torch
from torch import nn, istft
from torchaudio.transforms import Spectrogram, MelScale, InverseMelScale
import torch.nn.functional as F
import torch.autograd.profiler as profiler
from torch.autograd.profiler import record_function

class Audio_Model(nn.Module):</code></pre>
    This file has been truncated. <a href="https://gist.github.com/JuanFMontesinos/90db013f3c736a0a702faae28e73f4ff" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
This occur at least for RTX 3090, Quadro P6000, Titan V and some other gpus.</p>