How to make `torch.max` seach only few columns in the output of the model?

palguna_gopireddy: <p>I have designed a model with output nodes <em>(400 in my case)</em> not equal to the number of classes <em>(5 in my case)</em>. When I train this model, only first 0 to 4 columns in the output are non-zero, remaining 5 to 400 column values are zero.<br>
So when I used    <code>_, pred_labels = torch.max(y_predict, 1)</code>, it is producing predicted labels in between 0 to 4.<br>
That is how I am able to get the accuracy of 97%.</p>
<p>How can I make <code>torch.max</code> seach only <code>0 to 4 columns of each row of the model output</code>. If so, Does this valid as correct way of training?</p>

ptrblck: <p>You can slice the output via <code>y_predict[:, :5]</code> to use only the first 5 columns.</p>

palguna_gopireddy: <p>Thank you. Is this valid as correct way of training, if I put <code>y_predict[:,:5]</code> in the training after <code>y_predict = model(images.to(device))</code> in the training stage?</p>

ptrblck: <p>Yes, you would still waste compute to calculate the unused logits, but I don’t see other issues using this approach.</p>

palguna_gopireddy: <p>Thank you. I have tried putting output nodes of the model as 5 instead of 400. But It has effected the results. So if I want to reduce computation I think I have to change the whole architecture, which is again a burden. So I will slice the <code>y_predict</code>, as training stage is not violating anything.</p>
<p>Just one more information asking to gain knowledge, Does Pytorch have any <code>torch.nn</code> module or other torch module which gives this slicing or otherway to get this first 5 columns, so that I can put it in the module architecture itself?"</p>

ptrblck: <aside class="quote no-group" data-username="palguna_gopireddy" data-post="5" data-topic="187339">
<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/p/df705f/48.png" class="avatar"> palguna_gopireddy:</div>
<blockquote>
<p>I have tried putting output nodes of the model as 5 instead of 400. But It has effected the results.</p>
</blockquote>
</aside>
<p>I’m not sure why this should be the case as the unused parameters won’t receive any gradients as they were not used as seen in this small example:</p>
<pre data-code-wrap="python"><code class="lang-python">lin = nn.Linear(10, 10, bias=False)
x = torch.randn(1, 10)
out = lin(x)

out[:, :5].mean().backward()
print(lin.weight.grad)
# tensor([[ 0.0022, -0.0048, -0.0104, -0.1546,  0.1455,  0.1668,  0.0132,  0.2682,
#           0.0032, -0.0224],
#         [ 0.0022, -0.0048, -0.0104, -0.1546,  0.1455,  0.1668,  0.0132,  0.2682,
#           0.0032, -0.0224],
#         [ 0.0022, -0.0048, -0.0104, -0.1546,  0.1455,  0.1668,  0.0132,  0.2682,
#           0.0032, -0.0224],
#         [ 0.0022, -0.0048, -0.0104, -0.1546,  0.1455,  0.1668,  0.0132,  0.2682,
#           0.0032, -0.0224],
#         [ 0.0022, -0.0048, -0.0104, -0.1546,  0.1455,  0.1668,  0.0132,  0.2682,
#           0.0032, -0.0224],
#         [ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
#           0.0000, -0.0000],
#         [ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
#           0.0000, -0.0000],
#         [ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
#           0.0000, -0.0000],
#         [ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
#           0.0000, -0.0000],
#         [ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
#           0.0000, -0.0000]])
</code></pre>
<aside class="quote no-group" data-username="palguna_gopireddy" data-post="5" data-topic="187339">
<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/p/df705f/48.png" class="avatar"> palguna_gopireddy:</div>
<blockquote>
<p>Does Pytorch have any <code>torch.nn</code> module or other torch module which gives this slicing or otherway to get this first 5 columns, so that I can put it in the module architecture itself?"</p>
</blockquote>
</aside>
<p>You could write a custom <code>nn.Module</code> as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">class Slice(nn.Module):
    def __init__(self, index):
        super().__init__()
        self.index = index
        
    def forward(self, x):
        return x[:, :self.index]
    
model = nn.Sequential(
    nn.Linear(10, 10),
    Slice(5),
    nn.Linear(5, 20)
)

x = torch.randn(1, 10)
out = model(x)
print(out.shape)
# torch.Size([1, 20])
</code></pre>

palguna_gopireddy: <aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="187339">
<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’m not sure why this should be the case as the unused parameters won’t receive any gradients as they were not used as seen in this small example:</p>
</blockquote>
</aside>
<p>May be it’s because the last stage of the model is  <code>AdaptiveAvgPooling(1)</code>,then <code>Dropout</code> and <code>Flattening</code> layers. I have tried changing from 400 to 5 before this last stage, as I thought it is where I am producing extra 395.<br>
Can <code>AdaptiveAvgPooling</code>  reduce incoming 400 to 5 channels (with 1*1 feature maps ) or is there other pooling technique?</p>

ptrblck: <p>Your currently used pooling layer seems to reduce the spatial size of the activation to <code>1</code>, so I assume the <code>400</code> values are in the channel dimension? If so, you could use e.g. an <code>nn.Linear</code> layer to reduce the number of activations to <code>5</code>. Alternatively you could also use an <code>nn.Conv2d</code> layer with a kernel size of <code>1x1</code>, but this would be equivalent to the linear layer.</p>

palguna_gopireddy: <p>Yes, 400 values are in the channel dimension. Thank you.</p>