How to speedup the custom convolution process?

Loud_BoomBox: <p>Hi,<br>
I am trying to perform a custom convolution operation. I am using a c++ extension to perform the convolution operation using a lookup table. The trained model is an Encoder-Decoder architecture, with encoder as DenseNet-161.</p>
<p>Following are the input and weight dim for the 1st layer of decoder, over which I’m doing custom convolution op.</p>
<pre><code class="lang-auto">input dimension = [1,2208,15,20]
kernel dim = [1104,2208,1,1]
</code></pre>
<p>And this is the manner in which I reshape them and call the c++ extension for custom conv2d function:</p>
<pre><code class="lang-auto">img = F.pad(input= input, pad= (pad, pad, pad, pad), mode='constant', value= 0)

patches = img.unfold(2,kh,dh).unfold(3,kw,dw).reshape(batch_size, in_channels*kh*kw, -1)
result = torch.zeros(batch_size, out_channel, out_h, out_w)
kernel = kernel.reshape(out_channel, -1)

result = custom_conv.conv_two_d(patches, kernel, lookup_table, out_h, out_w)
</code></pre>
<p>The c++ implementation is given below:</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;torch/types.h&gt;
#include &lt;vector&gt;
#include &lt;THC/THC.h&gt;
#include &lt;iostream&gt;
#include &lt;ATen/NativeFunctions.h&gt;
#include &lt;ATen/Config.h&gt;
#include &lt;cuda_runtime_api.h&gt;

#include&lt;cmath&gt;
#include&lt;iostream&gt;
#include&lt;cstdlib&gt;

torch::Tensor conv_two_d(torch::Tensor input,
                           torch::Tensor weights,
                           torch::Tensor lookup_table,
                           int64_t outh, int64_t outw) {

    int64_t batch_size = input.size(0);
    int64_t out_channels = weights.size(0);
    torch::Tensor output = torch::zeros(torch::IntArrayRef({batch_size, out_channels, outh, outw})).cuda();
    torch::Tensor result = torch::zeros(torch::IntArrayRef({out_channels, outh*outw})).cuda();
    input = input.cuda();
    weights = weights.cuda();
    int sign=0;
    time_t start, end;
    std::cout&lt;&lt;"\n Starting..."&lt;&lt;std::endl;

    for (int elt=0; elt&lt;batch_size;elt++){
        torch::Tensor input_n = input[elt].cuda();
        for(int i=0; i&lt;weights.size(0); i++){
            start = time(NULL);

            for(int j=0; j&lt;input_n.size(1); j++){
                int64_t r=0;

                for(int k=0; k&lt;weights.size(1);k++){

                    float t1 = weights[i][k].item&lt;float&gt;();
                    float t2 = input_n[k][j].item&lt;float&gt;();
                    t1 = round(t1*1000); t2 = round(t2*1000);


                    if(t1&gt;255){t1=255;} if(t1&lt;-255){t1=-255;}                   
                    if(t2&gt;255){t2=255;} if(t2&lt;-255){t2=-255;}


                    if(t1&gt;0 &amp;&amp; t2&gt;0 || t1&lt;0 &amp;&amp; t2&lt;0){
                        t1=abs(t1); t2=abs(t2);
                        sign=1;
                    }

                    if(t1&lt;0 &amp;&amp; t2&gt;0 || t1&gt;0 &amp;&amp; t2&lt;0){
                        t1=abs(t1); t2=abs(t2);
                        sign=-1;
                    }


                    r+= lookup_table[t1][t2].item&lt;int&gt;()*sign;
                }

                result[i][j] = r/1000000;
            }

            end = time(NULL);
            std::cout&lt;&lt;"Time taken for iteration:"&lt;&lt;i&lt;&lt;" in batch:"&lt;&lt;elt&lt;&lt;" is:"&lt;&lt;end-start&lt;&lt;std::endl;
        }
        
        output[elt].add_(result.reshape(torch::IntArrayRef({out_channels, outh, outw})));
    }
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("conv_two_d", &amp;conv_two_d,"conv forward (CUDA)");
}
</code></pre>
<p>The time taken to complete each iteration is ~10-11 seconds… which in turn takes ~3.5 hrs to complete inference over a single image. (as shown below)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a182f36c11f2a9ae81c6d282553df07fd6a144c1.png" data-download-href="https://discuss.pytorch.org/uploads/default/a182f36c11f2a9ae81c6d282553df07fd6a144c1" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a182f36c11f2a9ae81c6d282553df07fd6a144c1.png" alt="image" data-base62-sha1="n2NqT4HdlwlfdjPKrk0PING06jf" width="690" height="191" data-dominant-color="8D7B86"><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">image</span><span class="informations">799×222 24.5 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>My doubts are:</p>
<p><em><strong>1) What can I possibly do to increase the speed at which the inference is done?</strong></em><br>
I have looked into the possibility of FFT_conv but I got to know that it’s mostly for larger kernels and images and saw it mentioned that, performing FFT_conv on smaller kernels might actually slow down the convolution process further.</p>
<p><strong><em>2) Is there a way in which I can reduce the inference time for smaller kernels</em>?</strong><br>
The model on which I am doing inference has only 1x1 and 3x3 kernels.</p>
<p>TIA</p>