Speedup F.unfold and F.fold

Zarek_Xu: <p>I am using a customized convolutional function, including F.unfold and F.fold. However, both training time and inference time is much longer than the original conv2d operation in pytorch. Anyone know how can I speed up F.unfold and F.fold? Thank you!</p>

albanD: <p>Hi,</p>
<p>unfold and fold should be very fast as they only play with stride in general. Are you sure they are responsible for the slowdown.<br>
If you compare with a general convolution algorithm, depending on the input, it is not always the most efficient to do the unfold, mm, fold.</p>

Zarek_Xu: <p>Hi,</p>
<p>Thank you so much for your kind reply! That is strange because I used torchprof to measure the time cost of each operation and found im2col (I believe it’s unfold) and col2img (fold) account for the most of the time cost. Below is the profiling results:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd_2_568x500.jpeg" alt="image" data-base62-sha1="wyLs1ooZWVNpMMq5uP5qUMVLw5L" width="568" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd_2_568x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd_2_852x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4339ded9cbdc5e018dc8ca689925bf0f1d1f2fd_2_1136x1000.jpeg 2x" data-dominant-color="3E1730"><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">1462×1285 388 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>

Zarek_Xu: <p>I also attach my unfold and fold code. I would be really appreciate if you can tell me where is the problem.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7daa125943ee6f27bed1281c631282adc4856a35.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7daa125943ee6f27bed1281c631282adc4856a35" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7daa125943ee6f27bed1281c631282adc4856a35_2_690x340.jpeg" alt="image" data-base62-sha1="hVG3LhKGIKKaxLFm3a9yHpMyded" width="690" height="340" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7daa125943ee6f27bed1281c631282adc4856a35_2_690x340.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7daa125943ee6f27bed1281c631282adc4856a35_2_1035x510.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7daa125943ee6f27bed1281c631282adc4856a35_2_1380x680.jpeg 2x" data-dominant-color="2C2E27"><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">1918×947 220 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>Thank you so much!</p>

Zarek_Xu: <p>If I use a general nn.conv2d, the training time of each batch (256 batch size) is only 80ms. However, if I use my customized conv2d, the training time will increase to 800ms.</p>

albanD: <p>That is indeed very surprising.<br>
Are all the timings measured with proper cuda synchronization?</p>
<p>Also could you share a text version of your module so that I can easily copy paste it and make an independent repro of the timing?</p>

Zarek_Xu: <p>Of course, thank you for your help. Below is the text code:</p>
<pre><code class="lang-python">class MyConv2d1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, sratio=0):
        super(MyConv2d1, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.sratio = sratio
        # print(self.sratio)
        self.weight = Parameter(torch.cuda.FloatTensor(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        h_in, w_in = input.shape[2:]
        h_out = h_in
        w_out = w_in
        # x: [batchsize ksize num_sliding]
        x = torch.nn.functional.unfold(input, kernel_size=self.kernel_size, padding=self.padding)

        batchsize = input.shape[0]
        ksize = self.in_channels*self.kernel_size*self.kernel_size
        num_sliding = x.shape[2]

        assert x.shape[1] == ksize


        w = self.weight
        out_unf = x.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
        out = torch.nn.functional.fold(out_unf, output_size=[h_out, w_out], kernel_size=1, padding=0, dilation=self.dilation, stride=1)

        return out
</code></pre>

albanD: <p>Running the following on CPU is slower but not that bad:</p>
<p>In particular, the fold is “almost free” as it does not do any copy.<br>
And the two other ops have a similar cost as they go through the big unfolded Tensor once each.</p>
<pre><code class="lang-python">import torch
from torch import nn
from torch.nn import Parameter
import time

class MyConv2d1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, sratio=0):
        super(MyConv2d1, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.sratio = sratio
        # print(self.sratio)
        # CHANGED: Remove cuda
        self.weight = Parameter(torch.rand(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        h_in, w_in = input.shape[2:]
        h_out = h_in
        w_out = w_in
        # x: [batchsize ksize num_sliding]
        start = time.time()
        x = torch.nn.functional.unfold(input, kernel_size=self.kernel_size, padding=self.padding)
        print(" Unfold: ", time.time() - start)

        batchsize = input.shape[0]
        ksize = self.in_channels*self.kernel_size*self.kernel_size
        num_sliding = x.shape[2]

        assert x.shape[1] == ksize


        w = self.weight
        start = time.time()
        out_unf = x.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
        print(" Mm: ", time.time() - start)
        start = time.time()
        out = torch.nn.functional.fold(out_unf, output_size=[h_out, w_out], kernel_size=1, padding=0, dilation=self.dilation, stride=1)
        print(" Fold: ", time.time() - start)

        return out

in_c = 10
out_c = 15
kernel = 3
padding = 1 # Custom Module only works when output is same size as input
inp = torch.rand(2, in_c, 50, 50)

conv = torch.nn.Conv2d(in_c, out_c, kernel, padding=padding)
custom_conv = MyConv2d1(in_c, out_c, kernel, padding=padding)

start = time.time()
out = conv(inp)
print("Regular: ", time.time() - start)

start = time.time()
out = custom_conv(inp)
print("Custom: ", time.time() - start)
</code></pre>

Zarek_Xu: <p>Sorry, I am a little bit confused. First, I think my code is running on the GPU. Second, I re-compare Myconv2d1 and nn.Conv2d. The training time of each batch is 450ms vs 100ms. I guess the extra 300ms is coming from fold and unfold, is that correct?<br>
Thank you!</p>

Zarek_Xu: <p>Is it that possible the the cost of these two operations during training are much larger than inference?</p>

Zarek_Xu: <p>I also compare the inference time cost which is: customize(200ms) vs conv2d (30ms). According to your testing results, the difference should not be such large, right? Sorry, I am a new learner, there are a lot of things that confused me. Thank you.</p>

albanD: <p>You are correct. I changed it to run on CPU because timing on GPU is a bit trickier (synchonization etc).</p>
<p>The overall backward time is similar to the forward one:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import nn
from torch.nn import Parameter
import time

class MyConv2d1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, sratio=0):
        super(MyConv2d1, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.sratio = sratio
        # print(self.sratio)
        # CHANGED: Remove cuda
        self.weight = Parameter(torch.rand(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        h_in, w_in = input.shape[2:]
        h_out = h_in
        w_out = w_in
        # x: [batchsize ksize num_sliding]
        start = time.time()
        x = torch.nn.functional.unfold(input, kernel_size=self.kernel_size, padding=self.padding)
        print(" Unfold: ", time.time() - start)

        batchsize = input.shape[0]
        ksize = self.in_channels*self.kernel_size*self.kernel_size
        num_sliding = x.shape[2]

        assert x.shape[1] == ksize


        w = self.weight
        start = time.time()
        mat_inp1 = x.transpose(1, 2)
        mat_inp2 = w.view(w.size(0), -1).t()
        mat_out = mat_inp1.matmul(mat_inp2)
        out_unf = mat_out.transpose(1, 2)
        print(" Mm: ", time.time() - start)
        start = time.time()
        out = torch.nn.functional.fold(out_unf, output_size=[h_out, w_out], kernel_size=1, padding=0, dilation=self.dilation, stride=1)
        print(" Fold: ", time.time() - start)


        # Setup backward hooks to time backward
        # This does not count the backward of all the view/transposes
        backward_times = []
        out.register_hook(lambda grad: backward_times.append(time.time()))
        mat_out.register_hook(lambda grad: backward_times.append(time.time()))
        x.register_hook(lambda grad: backward_times.append(time.time()))

        out_unf.register_hook(lambda grad: print(" Back fold: ", time.time() - backward_times[-1]))
        mat_inp1.register_hook(lambda grad: print(" Back mm: ", time.time() - backward_times[-1]))
        input.register_hook(lambda grad: print(" Back unfold: ", time.time() - backward_times[-1]))


        return out

in_c = 10
out_c = 15
kernel = 3
padding = 1 # Custom Module only works when output is same size as input
inp = torch.rand(2, in_c, 50, 50, requires_grad=True)

conv = torch.nn.Conv2d(in_c, out_c, kernel, padding=padding)
custom_conv = MyConv2d1(in_c, out_c, kernel, padding=padding)

start = time.time()
out = conv(inp)
print("Regular: ", time.time() - start)
start = time.time()
out.sum().backward()
print("Regular backward: ", time.time() - start)


start = time.time()
out = custom_conv(inp)
print("Custom: ", time.time() - start)
start = time.time()
out.sum().backward()
print("Custom backward: ", time.time() - start)
</code></pre>
<blockquote>
<p>I also compare the inference time cost which is: customize(200ms) vs conv2d (30ms). According to your testing results, the difference should not be such large, right?</p>
</blockquote>
<p>Right, at least not on CPU.</p>

Zarek_Xu: <p>Thank you so much for helping me! In that case, do you have ideas that can speed up fold and unfold on GPU? I want to use my customized conv2d to training models, but the time cost is really expensive.</p>

Zarek_Xu: <p>Interesting, when I increased the batchsize in your code (inp = torch.rand(2, in_c, 50, 50, requires_grad=True) → inp = torch.rand(512, in_c, 50, 50, requires_grad=True) ), the time cost of both training and inference of costumize conv2d increased dramatically. Am I correct? Thank you!</p>

albanD: <p>Both gets slower right?</p>
<pre><code class="lang-auto">Regular:  0.21726393699645996
Regular backward:  0.42275333404541016
 Unfold:  0.357952356338501
 Mm:  0.4160928726196289
 Fold:  0.08525800704956055
Custom:  0.8644123077392578
 Back fold:  0.07413196563720703
 Back mm:  0.2995262145996094
 Back unfold:  0.44156312942504883
Custom backward:  0.8373520374298096
</code></pre>
<p>But the custom is still between 2 and 4 times slower.</p>
<blockquote>
<p>Thank you so much for helping me! In that case, do you have ideas that can speed up fold and unfold on GPU?</p>
</blockquote>
<p>I am not sure if there is anything you can do.<br>
The thing is that on GPU, cudnn has very fine tuned algorithms to do convolution that are significantly better than MM.<br>
Can you try comparing the runtimes on GPU when cudnn is disabled via <code>torch.backends.cudnn.enabled = False</code>?</p>

Zarek_Xu: <p>You are right! When I set torch.backends.cudnn.enabled = False, the training time cost increased from 50ms to 200ms, which similar to customized conv2d.</p>

albanD: <p>Ok.<br>
So the unfold/mm/fold is as good as a specialized conv implementation based on mm<br>
But cudnn implementation is much better for your particular sizes.<br>
Sorry about that <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

Zarek_Xu: <p>You help me a lot! I really appreciate that. Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>