Convolution using FFT is slower than Conv2d :thinking:

username: eduardo4jesus
cooked: <p>Hello folks,</p>
<p>According to my understanding, a convolution operation requires, for a kernel of size k, input batch of size B, input channels c, output channels f, image size n by n.</p>
<p>$$B * f * c * k^2 * n^2$$</p>
<p>The convolution can also be done on the fft domain by a element wise multiplication. Disregarding the cost of the FFT operations, for simplicity, we would have (note each complex multiplication requires 4 real ones)</p>
<p>$$<br>
4 * B * f * c * n^2<br>
$$</p>
<p>ps.: don’t know why formulas in markdown here are not working</p>
<p>This plot show that using FFT should be faster, even when I add the cost for the FFT operations themselves.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/8/b8a6256fb50ca5a987c53da5f394f95544539350.png" data-download-href="https://discuss.pytorch.org/uploads/default/b8a6256fb50ca5a987c53da5f394f95544539350" title="Analysis_Operations"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8a6256fb50ca5a987c53da5f394f95544539350_2_690x437.png" alt="Analysis_Operations" data-base62-sha1="qltOUOm0SFfvgTVK8OnxA6zZAju" width="690" height="437" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8a6256fb50ca5a987c53da5f394f95544539350_2_690x437.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8a6256fb50ca5a987c53da5f394f95544539350_2_1035x655.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8a6256fb50ca5a987c53da5f394f95544539350_2_1380x874.png 2x" data-dominant-color="FBFCFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Analysis_Operations</span><span class="informations">4767×3024 299 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, when I tested it, see <a href="https://colab.research.google.com/gist/eduardo4jesus/c6ea31bd245843597517b0021675d5ed/sadness-conv-speed-test-torch.ipynb" rel="noopener nofollow ugc">Gist</a> the time used for the multiplication operations is away longer than anything else.</p>
<p>Also, I had to be very careful, anything change could easily cause allocation error due to insufficient memory on GPU.</p>
<p>Can anyone give me a hint on how Con2d is faster than the FFT version? It suppose to be the other way around. <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>

username: ezyang
cooked: <p>At least part of your performance loss is you had to reimplement complex multiplication by hand using an unfused series of operations. Fusing the operations would be a big win. You should be able to use the PyTorch JIT’s fuser to test this.</p>

username: eduardo4jesus
cooked: <p>Thanks for the feedback. I am not familiar with the term PyTorch JIT’s fuser. Is <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="nofollow noopener">this</a> where I should be reading from?</p>
<p>PS:<br>
I just got these two useful links. I will go through it and try a second implementation.</p>
<div class="lazyYT" data-youtube-id="St3gdHJzic0" data-youtube-title="Research to Production: PyTorch JIT/TorchScript Updates - Michael Suo" data-width="480" data-height="270" data-parameters="feature=oembed&amp;wmode=opaque"></div><br>
<a href="https://pytorch.org/docs/stable/jit.html" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/docs/stable/jit.html</a>

username: eduardo4jesus
cooked: <p>I wrapped my functions calls inside a <code>torch.nn.Module</code>, and used <code>jit</code></p>
<p>I know that <a href="https://colab.research.google.com/gist/eduardo4jesus/eff1fd82a0e57c15e333f67db4587a04/sadness-mnist-32-conv-speed-torch.ipynb" rel="nofollow noopener">my code</a> is really messy, didn’t have still time to make it more readable neither better written. But, hopefully, is good enough to see that the timers are still in favor of the traditional <code>Conv2d</code>.</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>Conv2d</th>
<th>Conv2dFFT</th>
<th>Conv2dFFT (without JIT)</th>
</tr>
</thead>
<tbody>
<tr>
<td>0.2027</td>
<td>0.8736</td>
<td>2.9177</td>
</tr>
</tbody>
</table>
</div><p>I am using <code>jit</code> like this, not sure how to use <code>fuser</code> as you mentioned.</p>
<pre><code class="lang-auto">class Conv2d(torch.nn.Module):
  def forward(self, x, filters):
    return torch.nn.functional.conv2d(x, filters, padding=2)

a = torch.rand(128, 64, 32, 32).cuda()
conv2d_spatial = Conv2d()
conv2d_spatial_opt = torch.jit.trace(conv2d_spatial, (a, filters)

...
.
.
.
...

class Conv2dFFT(torch.nn.Module):
  def forward(self, x, h):
    return fft_conv(x, h)

conv2d_fft = Conv2dFFT()
conv2d_fft_opt = torch.jit.trace(conv2d_fft, (a, filters))
</code></pre>