Differentiable Threhold for Conv2d output

username: Ajinkya_Ambatwar
cooked: <p>Hi I am trying to to implement learnable thresholding for the output of conv2d layers in the CNN.<br>
I want to make the output of conv layer to 0 for values between -threshold to +threshold. Also I want to differentiate the loss function, with respect to the threshold and update the threshold with each iteration.<br>
I have come across <code>.clamp </code> and the nn.threshold method in pytorch but I am unknown about about how to use them in class definition(eg. alexnet) and will autograd support this. If not what is the way around?</p>
<p>Thank you!!</p>

username: KFrank
cooked: <p>Hi Ajinkya!</p>
<aside class="quote no-group" data-username="Ajinkya_Ambatwar" data-post="1" data-topic="75345" data-full="true">
<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/ajinkya_ambatwar/48/11828_2.png" class="avatar"> Ajinkya_Ambatwar:</div>
<blockquote>
<p>I want to make the output of conv layer to 0 for values between -threshold to +threshold.</p>
</blockquote>
</aside>
<p>You want a differentiable, “soft” threshold function.  I think that<br>
<a href="https://pytorch.org/docs/stable/nn.html#torch.nn.Tanhshrink" rel="noopener nofollow ugc">nn.Tanhshrink</a> gives you most of what you want.  You can put in<br>
a threshold parameter like this:</p>
<pre><code class="lang-auto">threshold * nn.functional.tanhshrink (x / threshold)
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>Also I want to differentiate the loss function, with respect to the threshold and update the threshold with each iteration.</p>
</blockquote>
</aside>
<p>I’ve never done this, but I believe that if you make your<br>
<code>threshold</code> value a <code>nn.Parameter</code> and include it in your model,<br>
it should work.</p>
<p>The method in this post:</p>
<aside class="quote quote-modified" data-post="2" data-topic="50436">
  <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/akernel/48/12361_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-train-my-custom-parameter/50436/2">How to train my custom parameter?</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    You could write your own nn.Module that includes a trainable alpha parameter: 
class MyLayer(nn.Module):
    def __init__(self, *args, **kwargs):
        self.alpha = nn.Parameter(torch.tensor(0.))
        # Other necessary setup
        super().__init__(*args, **kwargs)

    def forward(self, x):
        # Necessary forward computations
        return output

Using this MyLayer module as part of a network architecture will allow you to learn alpha along with all of the other usual parameters.  …
  </blockquote>
</aside>

<p>looks plausible to me (but I can’t vouch for it).</p>
<p>Best.</p>
<p>K. Frank</p>

username: Ajinkya_Ambatwar
cooked: <p>Thank you Frank,<br>
I would look into the solution you are proposing</p>