Making prediction with argmax

hiepnguyen034: <p>I was following <a href="https://pytorch.org/tutorials/beginner/nn_tutorial.html" rel="nofollow noopener">this tutorial</a> and was confused by the way we make final prediction. Let’s say the model is defined as</p>
<pre><code class="lang-auto">loss_func = F.cross_entropy

def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, opt = get_model()
print(loss_func(model(xb), yb))

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))
</code></pre>
<p>. Then, the final prediction can be achieved by calling <code>torch.argmax(model(test_data), dim=1</code>). This means that <code>y_pred=model(test_data)</code> will return a vector that has a probability value for each digit. Then, my question is, how can PyTorch know index <code>0</code> of y_pred will correspond to the probability of the digit being <code>0</code>, index <code>1</code> will correspond to the probability of the digit being <code>1</code>, and so on? Since our <code>yb</code> is a one-dimensional vector that looks like <code>[5, 0, 4, ..., 8, 4, 8]</code> instead of a one-hot matrix, how can PyTorch know to link the index of <code>y_pred</code> with the correct label?</p>
<p>Thank you!</p>

KFrank: <p>Hi Hiep!</p>
<aside class="quote no-group quote-modified" data-username="hiepnguyen034" data-post="1" data-topic="49526" 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/hiepnguyen034/48/12872_2.png" class="avatar"> hiepnguyen034:</div>
<blockquote>
<p>…</p>
<pre><code class="lang-auto">loss_func = F.cross_entropy
</code></pre>
<p>…</p>
<p>Then, the final prediction can be achieved by calling <code>torch.argmax(model(test_data), dim=1</code>). This means that <code>y_pred=model(test_data)</code> will return a vector that has a probability value for each digit. Then, my question is, how can PyTorch know index <code>0</code> of y_pred will correspond to the probability of the digit being <code>0</code>, index <code>1</code> will correspond to the probability of the digit being <code>1</code>, and so on? Since our <code>yb</code> is a one-dimensional vector that looks like <code>[5, 0, 4,  ..., 8, 4, 8]</code> instead of a one-hot matrix, how can PyTorch know to link the index of <code>y_pred</code> with the correct label?</p>
<p>Thank you!</p>
</blockquote>
</aside>
<p>The short answer is that <code>nn.functional.cross_entropy</code><br>
one-hots your class labels for you.</p>
<p>The number-one rule is that the output of your network means<br>
whatever you train it to mean.</p>
<p>More directly to your question:</p>
<p>You are using <code>nn.functional.cross_entropy</code> as your loss<br>
function.  Your <code>yb</code> are integer class labels, one per sample.</p>
<p>Conceptually, implicitly under the hood, <code>cross_entropy</code><br>
one-hots your <code>yb</code>, implicitly softmaxes your <code>pred</code>, and then<br>
implicitly calculates the cross-entropy of “softmax (pred)” and<br>
“one-hot (yb)”.</p>
<p>The output of your model (the <code>y_pred</code>) should be understood<br>
as logits.  They (implicitly) get turned into probabilities when<br>
<code>cross_entropy</code> (implicitly) softmaxes them.  So you are training<br>
your model to output (for each sample) a vector of length nClass,<br>
where the value for index i is the logit (sort of like the probability)<br>
of that sample being of class i.  (Finally, you take the <code>argmax</code> of<br>
your prediction vector.  This finds the index of the logit with the<br>
largest value – that is the index that your model predicts as having<br>
the highest probability of being the class label, and you take this<br>
as being the predicted class label.)</p>
<p>“PyTorch knows to link the index of  <code>y_pred</code>  with the correct label”<br>
because you trained your network to do so.</p>
<p>Good Luck.</p>
<p>K. Frank</p>