Depthwise Separable Convolution: 7x Fewer Parameters, But Only 1.55x Speedup?

anonymous_anonymous1: <p>Hi everyone,</p>
<p>I’ve implemented and benchmarked <strong>Depthwise Separable Convolutions (DWSConv)</strong> against standard convolutions to compare their performance on a GPU using PyTorch. I’m seeking feedback on both my implementation and the relevance of my benchmark.</p>
<p>Here’s my code for both layers:</p>
<pre><code class="lang-auto">from time import time

import torch
from torch import nn


class Conv(nn.Module):
    """Standard convolution"""

    def __init__(self, cin, cout, k, s, p):
        """Initialize Conv layer with given arguments including activation."""
        super().__init__()
        self.conv = nn.Conv2d(cin, cout, k, s, p, groups=1, bias=False)
        # No BatchNorm2d because one can fuse it with conv2d after training
        self.act = nn.ReLU()

    def forward(self, x):
        return self.act(self.conv(x))


class DWSConv(nn.Module):
    """DepthWise Separable Conv =  Depthwise Conv + Pointwise Conv"""

    def __init__(self, cin, cout, k, s, p):
        """Initialize Conv layer with given arguments including activation."""
        super().__init__()
        self.dw_conv = nn.Conv2d(cin, cin, k, s, p, groups=cin, bias=False) # Depthwise layer: cout=cin + groups=cin
        # No BatchNorm2d because one can fuse it with conv2d after training
        self.act_dw = nn.ReLU()
        self.pw_conv = nn.Conv2d(cin, cout, 1, 1, 0, groups=1, bias=False)  # Pointwise layer: k=1, s=1, p=0
        # No BatchNorm2d because one can fuse it with conv2d after training
        self.act_pw = nn.ReLU()

    def forward(self, x):
        """Apply convolution, batch normalization and activation to input tensor."""
        return self.act_pw(self.pw_conv(self.act_dw(self.dw_conv(x))))
    

device = "cuda"
cin, cout, k, s, p = 16, 32, 3, 2, 1

x = torch.randn(1024, cin, 320, 320).to(device).half()

conv_layer = Conv(cin, cout, k, s, p).to(device).half()
dwsconv_layer = DWSConv(cin, cout, k, s, p).to(device).half()


start = time()
_ = dwsconv_layer(x)
torch.cuda.synchronize()
print(f"(WARMUP) Duration for the DWS conv layer: {(time()-start)*1e3:.2f}ms")
start = time()

start = time()
_ = dwsconv_layer(x)
torch.cuda.synchronize()
print(f"Duration for the DWS conv layer: {(time()-start)*1e3:.2f}ms")


start = time()
_ = conv_layer(x)
torch.cuda.synchronize()
print(f"(WARMUP) Duration for the classical conv layer: {(time()-start)*1e3:.2f}ms")
start = time()

start = time()
_ = conv_layer(x)
torch.cuda.synchronize()
print(f"Duration for the classical conv layer: {(time()-start)*1e3:.2f}ms")
start = time()

print(f"Number of weights in classical conv: {conv_layer.conv.weight.nelement()}")
print(f"Number of weights in DWS conv: {dwsconv_layer.dw_conv.weight.nelement() + dwsconv_layer.pw_conv.weight.nelement()}")
</code></pre>
<h3><a name="p-457278-results-1" class="anchor" href="#p-457278-results-1"></a>Results:</h3>
<ul>
<li><strong>Depthwise Separable Convolution (DWSConv)</strong>:
<ul>
<li>Execution time: <strong>20.65 ms</strong></li>
<li>Number of parameters: <strong>656</strong></li>
</ul>
</li>
<li><strong>Standard Convolution</strong>:
<ul>
<li>Execution time: <strong>31.13 ms</strong></li>
<li>Number of parameters: <strong>4608</strong></li>
</ul>
</li>
</ul>
<h3><a name="p-457278-the-puzzle-2" class="anchor" href="#p-457278-the-puzzle-2"></a>The Puzzle:</h3>
<p>DWSConv has <strong>7x fewer parameters</strong> (656 vs 4608), yet it only gives a <strong>1.55x speedup</strong> (20.65 ms vs 31.13 ms).</p>
<h3><a name="p-457278-additional-issue-with-larger-inputs-3" class="anchor" href="#p-457278-additional-issue-with-larger-inputs-3"></a>Additional Issue with Larger Inputs:</h3>
<p>When I use larger input sizes like this:</p>
<pre><code class="lang-auto">cin, cout, k, s, p = 16, 32, 3, 2, 1
x = torch.randn(19_000, cin, 64, 128).to(device).half()
</code></pre>
<p>The standard convolution processes it without any issue, but the DWSConv throws this error:</p>
<pre><code class="lang-auto">RuntimeError: Expected canUse32BitIndexMath(input) &amp;&amp; canUse32BitIndexMath(output) to be true, but got false. 
(Could this error message be improved? If so, please report an enhancement request to PyTorch.)
</code></pre>
<p>This suggests that <strong>intermediate tensors</strong> in DWSConv could exceed the indexing limit of <code>2^31</code> elements. This is puzzling, especially since the standard <code>Conv2d</code> should handle more elements but doesn’t encounter this issue.</p>
<h3><a name="p-457278-my-question-4" class="anchor" href="#p-457278-my-question-4"></a>My Question:</h3>
<ol>
<li>Why is the speedup much smaller compared to the reduction in parameters?</li>
<li>Why does DWSConv hit an indexing limitation with large inputs while Conv2d does not?</li>
</ol>
<p>Looking forward to your insights!</p>