Bypass individual channels in nn.Conv2d

maaft: <p>In a recent paper <a href="https://arxiv.org/abs/2003.11883" class="inline-onebox" rel="noopener nofollow ugc">[2003.11883] DCNAS: Densely Connected Neural Architecture Search for Semantic Image Segmentation</a> the authors are talking about forwarding only 25% of random sampled channels of a featuremap through a conv-layer while “bypassing” the others to reduce the memory footprint.</p>
<p>Here is the corresponding section in that paper:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2ab01fb78fa84a2e5899a400a524db60f13c38fb.png" data-download-href="https://discuss.pytorch.org/uploads/default/2ab01fb78fa84a2e5899a400a524db60f13c38fb" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2ab01fb78fa84a2e5899a400a524db60f13c38fb.png" alt="image" data-base62-sha1="65DoDQxoqZwVXfshXSCHpoVEE9t" width="690" height="392" data-dominant-color="ECEBEB"><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">755×430 76.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>How would you achieve something like this in PyTorch?</p>
<p>I figured you could achieve this with something like in the following example, but maybe there is another simpler and obvious way that I am missing?</p>
<pre data-code-wrap="python"><code class="lang-python">class SampleConv(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
        super().__init__(
            in_channels, out_channels,
            kernel_size, stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

        self.channel_sample_ratio = 0.25

    def forward(self, input):
        sampled_channels = torch.rand(x.size(1), device=x.device, requires_grad=False)

        weight = self.weight[:, sampled_channels &lt; self.channel_sample_ratio]
        bias = self.bias

        y = F.conv2d(input[:, sampled_channels &lt; self.channel_sample_ratio], weight, bias, self.stride, self.padding, self.dilation, self.groups)
        return y
</code></pre>
<p>Also I measured the memory consumption of above <code>SampleConv</code> and its regular nn.Conv2d-equivalent and it doen’t seem like that it requires 75% less memory.</p>

ptrblck: <p>I’m not sure I understand the section correctly, but it seems that 25% of the input channels is randomly sampled and passed to the conv layer (or "mixed transformation of <code>O</code>?), while the remaining channels are “bypassed”?</p>
<p>Does the second part mean that 75% of the input channels are just concatenated to the output of the conv operaration?</p>

maaft: <p>I think you can assume that <code>input_channels = output_channels</code> so the addition in eq. 7 could make sense (instead of concatenation).</p>
<p>But apart from that, I understood the section in the same way as you.</p>
<p>Disabling gradients for certain channels is not possible, correct?</p>

ptrblck: <p>The number of output channels is defined by the number of kernels.<br>
If you want to remove 75% of the input channels and would like to keep the same amount of output channels, you would just remove kernels.</p>
<p>I might be wrong, but that’s not my understanding of the section.</p>
<aside class="quote no-group" data-username="maaft" data-post="3" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/4af34b/48.png" class="avatar"> maaft:</div>
<blockquote>
<p>Disabling gradients for certain channels is not possible, correct?</p>
</blockquote>
</aside>
<p>You could zero out the gradients for specific channels.</p>

maaft: <p>So you would also do something like my <code>SampleConv</code> approach from my first post i.e. selecting certain kernels? Would it make sense to also select only output channels?</p>
<p>And yes, at second glance I think you are totally right. Addition makes no sense at all so concatenation might be the way to go.</p>
<p>Unfortunately I still can’t wrap my head around their claimed 75% decrease in memory consumption.</p>

ptrblck: <aside class="quote no-group" data-username="maaft" data-post="5" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/4af34b/48.png" class="avatar"> maaft:</div>
<blockquote>
<p>So you would also do something like my <code>SampleConv</code> approach from my first post i.e. selecting certain kernels?</p>
</blockquote>
</aside>
<p>In your code you are slicing the input channels, i.e. you are making the kernels “smaller”:</p>
<pre data-code-wrap="python"><code class="lang-python">self.weight[:, sampled_channels &lt; self.channel_sample_ratio]
</code></pre>
<p>The <code>weight</code> parameter is defined as <code>[out_channels, in_channels, height, width]</code>.</p>
<aside class="quote no-group" data-username="maaft" data-post="5" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/4af34b/48.png" class="avatar"> maaft:</div>
<blockquote>
<p>Would it make sense to also select only output channels?</p>
</blockquote>
</aside>
<p>I don’t understand the posted section as removing kernels, but the input channels and bypassing the removed channels (i.e. not applying the convolution on them). But I’m still unsure. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="maaft" data-post="5" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/4af34b/48.png" class="avatar"> maaft:</div>
<blockquote>
<p>Unfortunately I still can’t wrap my head around their claimed 75% decrease in memory consumption.</p>
</blockquote>
</aside>
<p>That’s also, why I’m thinking, that I misunderstand the section, so let’s better wait for smarter users who could help us out. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

maaft: <aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>That’s also, why I’m thinking, that I misunderstand the section, so let’s better wait for smarter users who could help us out. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"> Yep, lets wait for them. I also mailed the authors so maybe they will answer me.</p>

ptrblck: <p>Please update this thread once you get an answer, as I’m curious how it should be implemented.</p>

maaft: <p>Unfortunately, I didn’t receive an answer yet. In general it’s really a pity that in academia it is so hard to get an answer from the authors. /endofrant</p>
<p>Looking at fig 1. from that paper I am now sure that concatenation is meant:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2abaed06116c75fc09518e3957ba685ce18167ff.png" data-download-href="https://discuss.pytorch.org/uploads/default/2abaed06116c75fc09518e3957ba685ce18167ff" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2abaed06116c75fc09518e3957ba685ce18167ff_2_690x184.png" alt="image" data-base62-sha1="660xw3i6QzkIHOXU7im69xelgL5" width="690" height="184" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2abaed06116c75fc09518e3957ba685ce18167ff_2_690x184.png, https://discuss.pytorch.org/uploads/default/original/3X/2/a/2abaed06116c75fc09518e3957ba685ce18167ff.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/a/2abaed06116c75fc09518e3957ba685ce18167ff.png 2x" data-dominant-color="E7E8EC"><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">930×248 51.2 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>Those 3x3, 5x5, 7x7 sep. convs are operations from their NAS search space. So as you can see, in this example two feature maps are sampled from the input and fed through those operations while the other ones are bypassed and concatenated to the output.</p>
<p>Therefore, I came up with this “solution”:</p>
<pre data-code-wrap="python"><code class="lang-python">class SampleConv(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, sampling_ratio=0.25, stride=1, padding=0, dilation=1,
                 groups=1, bias=True):
        super().__init__(
            in_channels, out_channels,
            kernel_size, stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

        self.channel_sample_ratio = sampling_ratio

    def forward(self, x, sample=True):
        if sample:
            sampled_channels = torch.rand(x.size(1), device=x.device, requires_grad=False)

            weight = self.weight[:, sampled_channels &lt; self.channel_sample_ratio]
            weight = weight[sampled_channels &lt; self.channel_sample_ratio]
            if self.bias is not None:
                bias = self.bias[sampled_channels &lt; self.channel_sample_ratio]
            else:
                bias = None

            x[:, sampled_channels &lt; self.channel_sample_ratio] = F.conv2d(
                x[:, sampled_channels &lt; self.channel_sample_ratio], weight, bias, self.stride, self.padding,
                self.dilation, self.groups)
        else:
            x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return x

</code></pre>
<p>I measured the GPU memory consumption and it requires 50%-60% less memory as the regular <code>nn.Conv2d</code> counterpart. (Which is still a far cry from the stated 75%).</p>
<p>Also the L1-Error on estimating a <code>torch.randn</code> tensor is slightly higher (0.83 vs 0.799) on convergence than with a regular <code>nn.Conv2d</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">sample_conv = SampleConv(320, 320, 1, sampling_ratio=0.25).cuda()
normal_conv = nn.Conv2d(320, 320, 1).cuda()

optim_sample_conv = torch.optim.SGD(sample_conv.parameters(), lr=0.1, momentum=0.9)
optim_normal_conv = torch.optim.SGD(normal_conv.parameters(), lr=0.1, momentum=0.9)
target = torch.randn(1, 320, 48, 48).cuda()

for e in range(10000):
    # sample conv
    input_t = torch.ones(1, 320, 48, 48).cuda()
    y = sample_conv(input_t)
    loss = (y - target).abs().mean()
    optim_sample_conv.zero_grad()
    loss.backward()
    optim_sample_conv.step()

    with torch.no_grad():
        y = sample_conv(input_t, sample=False)
        l1_error = (y - target).abs().mean()

    print("SampleConv: Loss: {} L1-Error: {}".format(loss, l1_error))

    # regular nn.Conv2d
    input_t = torch.ones(1, 320, 48, 48).cuda()
    y = normal_conv(input_t)
    loss = (y - target).abs().mean()
    optim_normal_conv.zero_grad()
    loss.backward()
    optim_normal_conv.step()

    with torch.no_grad():
        y = normal_conv(input_t)
        l1_error = (y - target).abs().mean()

    print("nn.Conv2d: Loss: {} L1-Error: {}".format(loss, l1_error))
</code></pre>
<p>What do you think?</p>

ptrblck: <p>Yeah, the figure explains it pretty well. Thanks for sharing.</p>
<p>I think the general idea looks alright, and I assume you are using grouped convolutions, as you are slicing the <code>weight</code> in two dimensions?</p>

maaft: <aside class="quote no-group" data-username="ptrblck" data-post="10" data-topic="77021">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>I think the general idea looks alright, and I assume you are using grouped convolutions, as you are slicing the <code>weight</code> in two dimensions?</p>
</blockquote>
</aside>
<p>No, I didn’t use grouped convolutions yet. What would be the motivation behind this?</p>
<p>Another thing that I find confusing is that they use InvertedResidual Blocks (MobileNetV3) which look like <code>input_channels --&gt; hidden_channels --&gt; hidden_channels --&gt; output_channels</code>. But my implementation from above requires <code>input_channels == output_channels</code> so I can’t actually use it for the MBConv-Block (<code>input_channels != hidden_channels</code>).</p>
<p>Edit: I think we should only use this <code>SampleConv</code> on the middle conv (<code>hidden -&gt; hidden</code>) of the InvertedResidual Block. The in- and out-convolutions are just 1x1 and not that expensive.</p>
<p>From the figure it also seems like that the sampled channels are only concatenated and not “inplace-modified” like in my implementation. Wouldn’t that be very bad for succeeding layers as they see different types of feature-maps (produced with different kernels) at the same position?</p>