How to easily modify relu layer in pytorch?

standshik: <p>I am new to pytorch. I am trying to create a new activation layer, let’s call it topk,  that would work as follows. It will take a vector x of size n as input (result of multiplying previous layer output by weight matrix and adding bias) and  a positive integer k and would output a vector topk(x) of size n whose elements are<br>
(topk(x))_i =x_i if x_i is one of the top k elements of x, 0 otherwise.</p>
<p>While calculating gradient of topk(x), top k elements of x should have gradient 1, everything else 0.</p>
<p>How should I implement this? Can you please provide some sample code?</p>

KFrank: <p>Hi Standshik!</p>
<aside class="quote no-group" data-username="standshik" data-post="1" data-topic="118206" data-full="true">
<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/s/47e85d/48.png" class="avatar"> standshik:</div>
<blockquote>
<p>(topk(x))_i =x_i if x_i is one of the top k elements of x, 0 otherwise.</p>
</blockquote>
</aside>
<p>You can do this using <code>torch.topk()</code> and <code>.scatter_()</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; x = torch.randperm (10).float()
&gt;&gt;&gt; x.requires_grad = True
&gt;&gt;&gt; x_topk_val_ind = torch.topk (x, 5)
&gt;&gt;&gt; x_topk_val_ind
torch.return_types.topk(
values=tensor([9., 8., 7., 6., 5.], grad_fn=&lt;TopkBackward&gt;),
indices=tensor([8, 6, 3, 5, 1]))
&gt;&gt;&gt; topk = torch.zeros_like (x).scatter_ (0, x_topk_val_ind[1], x_topk_val_ind[0])
&gt;&gt;&gt; topk
tensor([0., 5., 0., 7., 0., 6., 8., 0., 9., 0.], grad_fn=&lt;ScatterBackward0&gt;)
&gt;&gt;&gt; topk.sum().backward()
&gt;&gt;&gt; x.grad
tensor([0., 1., 0., 1., 0., 1., 1., 0., 1., 0.])
&gt;&gt;&gt; x
tensor([4., 5., 3., 7., 2., 6., 8., 0., 9., 1.], requires_grad=True)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

standshik: <p>Hi Frank,<br>
Thank you so much. That was really helpful. My goal is use this as an activation in building a neural network. Suppose I define a simple network like the following where I use nn.ReLU() activation. How do I package the idea of topk together so that I can create ‘mypackage’, import it and use mypackage.topk(k) instead of nn.ReLU()?</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, input_size, hidden_size, out_size):
        super(Net, self).__init__()                    
        self.fc1 = nn.Linear(input_size, hidden_size)                              
        self.fc2 = nn.Linear(hidden_size, out_size)
        self.relu = nn.ReLU()

    def forward(self, x):                          
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out
</code></pre>

KFrank: <p>Hi Standshik!</p>
<aside class="quote no-group" data-username="standshik" data-post="4" data-topic="118206" data-full="true">
<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/s/47e85d/48.png" class="avatar"> standshik:</div>
<blockquote>
<p>How do I package the idea of topk together so that I can create ‘mypackage’, import it and use mypackage.topk(k) instead of nn.ReLU()?</p>
</blockquote>
</aside>
<p>If it were me, I wouldn’t bother packaging it – I would just define <code>topk</code><br>
as a function (rather than a class).  You could define it (either as a<br>
function or a class) in a separate package and import it (but how to<br>
do that is a python question, rather than specific to pytorch).</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">    def forward(self, x):                          
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out
</code></pre>
</blockquote>
</aside>
<p>No need to instantiate a <code>self.relu()</code> replacement – just define your<br>
<code>topk</code> function and call it in <code>forward()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self, x):                          
        out = self.fc1(x)
        # out = self.relu(out)
        out = my_topk_function (out, my_value_for_k)
        out = self.fc2(out)
        return out
</code></pre>
<p>Best.</p>
<p>K. Frank</p>