Using optimised depthwise convolutions

swibe: <p>Hi all,</p>
<p>Following <a href="https://github.com/pytorch/pytorch/pull/3057" rel="nofollow noopener">#3057</a> and <a href="https://github.com/pytorch/pytorch/pull/3265" rel="nofollow noopener">#3265</a>, I was excited to try out depthwise separable convolutions, but I’m having a hard time activating these optimised code paths. I’m currently getting no speedup over default convolutions.</p>
<p>Here are the two layer types that make up the bulk of my network:</p>
<pre><code># Depthwise
nn.Conv2d(in_chans, in_chans * k, kernel_size, groups = in_chans)

# Normal
nn.Conv2d(in_chans * k, out_chans, 1)
</code></pre>
<p>If I profile the network’s execution, I get the following (trimmed):</p>
<pre><code>-------------------------  ------------  ------------  ------------
Name                           CPU time     CUDA time         Calls
-------------------------  ------------  ------------  ------------
conv2d                        130.737us    1016.190us            14
cudnn_convolution             160.005us     843.504us             8
thnn_conv_depthwise2d          58.475us    1246.438us             6
</code></pre>
<p>My concern is that the depthwise convolutions are being handled by THNN, and not THCUNN, where the new optimisations are. I have a second network which replaces normal convolution with multiple depthwise convolutions, and runs much slower, despite performing much less computation.</p>
<p>Am I missing something obvious?</p>
<pre><code>&gt;&gt;&gt; torch.__version__
'0.4.0a0+82e995e'
&gt;&gt;&gt; torch.version.cuda
'8.0.61'
&gt;&gt;&gt; torch.backends.cudnn.version()
7005</code></pre>

smth: <p>are you profiling correctly?</p>
<p>If you are using python timing calls, you have to insert cuda synchronize calls to get correct timing:</p>
<pre><code class="lang-auto">torch.cuda.synchronize()
now = time.time()
# my cuda calls
torch.cuda.synchronize()
total_time = time.time() - now
</code></pre>

swibe: <p>I’m using the following, from the docs:</p>
<pre><code>with torch.cuda.profiler.profile():
    model(input_var) # Warmup CUDA memory allocator and profiler
    with torch.autograd.profiler.emit_nvtx():
        output = model(input_var)</code></pre>

smth: <p>hmmm, okay. what are your layer and input sizes?</p>

swibe: <p>I have those two layers wrapped up to create a drop-in replacement for <code>nn.Conv2d</code>, and they’re used in CIFAR-10 and ImageNet scale networks. The particular network that I profiled looks like this:</p>
<pre><code>Input: 256x3x32x32

Conv2d(  3, 128, kernel_size=3, padding=1)

Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)

Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)
        
MaxPool2d(kernel_size=3, stride=2)
# Channels now sized 15x15
        
Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)

Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)
        
MaxPool2d(kernel_size=3, stride=2)
# Channels now sized 7x7
        
Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)

Conv2d(128, 128, kernel_size=3, padding=1, groups=128)
Conv2d(128, 128, kernel_size=1, padding=0)
        
AvgPool2d(kernel_size=7)
# Channels now sized 1x1
        
Conv2d(128,  10, kernel_size=1, padding=0)
</code></pre>
<p>I’ve also tried the wrapped layers in copies of AlexNet, SqueezeNet and ResNet from the <code>torchvision</code> module. I haven’t profiled those. All trained to a reasonable accuracy, but each iteration was slower than their default counterparts.</p>

swibe: <p>Here’s a simple test script I wrote which executes a single conv layer.</p>
<p>I had to drop back to PyTorch v0.3 so I no longer see which backend is active, but I haven’t yet found a configuration where grouped convolution is markedly faster than ordinary convolution. e.g.:</p>
<pre><code>$ python3 group_conv.py --batch=256 --img-size=50 --in-channels=256 --out-channels=256 --kernel-size=3 --groups=True
---------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                    CPU time        CUDA time            Calls        CPU total       CUDA total
---------------  ---------------  ---------------  ---------------  ---------------  ---------------
ConvForward             77.271us      16688.654us                1         77.271us      16688.654us

$ python3 group_conv.py --batch=256 --img-size=50 --in-channels=256 --out-channels=256 --kernel-size=3 --groups=False
---------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                    CPU time        CUDA time            Calls        CPU total       CUDA total
---------------  ---------------  ---------------  ---------------  ---------------  ---------------
ConvForward             72.029us      15207.838us                1         72.029us      15207.838us
</code></pre>
<aside class="onebox githubgist">
  <header class="source">
      <a href="https://gist.github.com/db434/74fcfcf49d52aa08c2335ee916bac9f4" target="_blank" rel="nofollow noopener">gist.github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://gist.github.com/db434/74fcfcf49d52aa08c2335ee916bac9f4" target="_blank" rel="nofollow noopener">https://gist.github.com/db434/74fcfcf49d52aa08c2335ee916bac9f4</a></h4>
<h5>group_conv.py</h5>
<pre><code class="Python">import argparse
import os
import subprocess
import sys
import torch

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch", type=int, default=256)
    parser.add_argument("--img-size", type=int, default=32)</code></pre>
This file has been truncated. <a href="https://gist.github.com/db434/74fcfcf49d52aa08c2335ee916bac9f4" target="_blank" rel="nofollow noopener">show original</a>

<p>
</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Ziju_Feng: <p>I’ve seen similar results when working with 3D ResNext models where group convolution is used extensively (in group size of 32). In 3D, the ResNext model actually only has around 50% of parameters (and thus flops) compared to the Resnet counterpart but the forward/backward pass is 25%/60% slower with batch size 1 and 3% and 10% slower with batch size 128 split to 4 GPUs.</p>

swibe: <p>I’ve had some offline discussions about this. Here are my findings:</p>
<ul>
<li>
<code>thnn_conv_depthwise2d</code>, the internal function called, <strong>does</strong> use CUDA, but not CuDNN.</li>
<li>CuDNN 7’s implementation of grouped/depthwise convolution is up to 3x quicker in the forward pass, but always slower in the backward pass.</li>
<li>Choosing when to use CuDNN and when not to is very difficult to describe in a maintainable way. It will involve looking at many of the parameters of the conv layer, the current version of CuDNN, and is probably dependent on the GPU used as well.</li>
</ul>
<p>I think it’s probably a case of waiting until CUDA/CuDNN provide consistent benefits in more situations.</p>

smth: <p>nvidia said they will integrate some dedicated kernels similar to pytorch’s group conv into the next cudnn version and choose it when it’s faster.</p>

Ziju_Feng: <p>Would the next version of cudnn mean a major/minor version update like 8.0/7.1 or a patch level update like 7.0.6? Did Nvidia give us an estimate on the release date? Thanks.</p>

yaysummeriscoming: <p>Hello, based on the nvidia releases notes, some depthwise seperable convolution improvements have now made it into cudnn:</p>
<p><a href="https://docs.nvidia.com/deeplearning/sdk/cudnn-release-notes/rel_730.html#rel_730" class="onebox" target="_blank" rel="nofollow noopener">https://docs.nvidia.com/deeplearning/sdk/cudnn-release-notes/rel_730.html#rel_730</a></p>
<p>Is this now integrated with pytorch?</p>

yaysummeriscoming: <p><a class="mention" href="/u/smth">@smth</a>, <a class="mention" href="/u/ziju_feng">@Ziju_Feng</a>?  I had a look at the source code &amp; it seems the THNN implementation is still being used?</p>

smth: <p><a class="mention" href="/u/yaysummeriscoming">@yaysummeriscoming</a> I’ve checked the code as well, and yes even in the case of depthwise we are still dispatching to thnn depthwise kernels: <a href="https://github.com/pytorch/pytorch/blob/51f1c4fea5e066e3cbb7a7251114de4381d2a926/aten/src/ATen/native/Convolution.cpp#L330" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/51f1c4fea5e066e3cbb7a7251114de4381d2a926/aten/src/ATen/native/Convolution.cpp#L330</a></p>
<p>We haven’t used CuDNN implementation because those kernels aren’t very optimized according to an NVIDIA employee: <a href="https://github.com/pytorch/pytorch/issues/1708#issuecomment-384015108" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/1708#issuecomment-384015108</a></p>
<p>Once the situation changes and CuDNN has faster and optimized convolution routines for Depthwise, we will switch to dispatching to CuDNN.</p>
<p>I’ve opened and issue to track this: <a href="https://github.com/pytorch/pytorch/issues/15513" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/15513</a></p>

euler16: <p>Even when I compare the performance of depthwise separable conv3d vs the normal conv3d without gpu, I find that the latter still outperforms the former.</p>

tstandley: <p>I think the problem might be memory bandwidth.</p>
<p>if you implement Depthwise Separable convolution something like this:</p>
<pre><code class="lang-auto">class SeparableConv2d(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False):
        super(SeparableConv2d,self).__init__()

        self.conv1 = nn.Conv2d(in_channels,in_channels,kernel_size,stride,padding,dilation,groups=in_channels,bias=bias)
        self.pointwise = nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias)
    
    def forward(self,x):
        x = self.conv1(x)
        x = self.pointwise(x)
        return x
</code></pre>
<p>Then your layer winds up needing to write to twice as much memory compared to a typical conv2d.</p>
<p>This memory usage shouldn’t be necessary though. If the code was written as a single layer in CUDA, each pixel of the task could be done together and only the final result would need to be written to vram. This would likely result in a massive speedup for this operation.</p>