How to read the predicted label of a Neural Network with Cross Entropy Loss?
I am using a neural network to predict the quality of the Red Wine dataset, available on UCI machine Learning, using Pytorch, and Cross-Entropy Loss as the loss function.
This is my code:
input_size = len(input_columns)
hidden_size = 12
output_size = 6 #because there are 6 classes

#Loss function
loss_fn = F.cross_entropy

class WineQuality(nn.Module):
    def __init__(self):
        super().__init__()
        # input to hidden layer
        self.linear1 = nn.Linear(input_size, hidden_size)
        # hidden layer and output
        self.linear2 = nn.Linear(hidden_size, output_size)
        
    def forward(self, xb): 
        out = self.linear1(xb)
        out = F.relu(out)
        out = self.linear2(out)
        return out
    
    def training_step(self, batch):
        inputs, targets = batch 
        # Generate predictions
        out = self(inputs) 
        # Calcuate loss
        loss = loss_fn(out,torch.argmax(targets, dim=1))
        return loss
    
    def validation_step(self, batch):
        inputs, targets = batch
        # Generate predictions
        out = self(inputs)
        # Calculate loss
        loss = loss_fn(out, torch.argmax(targets, dim=1))
        return {'val_loss': loss.detach()}
        
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        return {'val_loss': epoch_loss.item()}
    
    def epoch_end(self, epoch, result, num_epochs):
        # Print result every 100th epoch
        if (epoch+1) % 100 == 0 or epoch == num_epochs-1:
            print(&quot;Epoch [{}], val_loss: {:.4f}&quot;.format(epoch+1, result['val_loss']))

model = WineQuality()

def evaluate(model, val_loader):
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)

def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    history = []
    optimizer = opt_func(model.parameters(), lr)
    for epoch in range(epochs):
        # Training Phase 
        for batch in train_loader:
            loss = model.training_step(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        # Validation phase
        result = evaluate(model, val_loader)
        model.epoch_end(epoch, result, epochs)
        history.append(result)
    return history

loss_value = evaluate(model, valid_dl)

#model=WineQuality()
epochs = 1000
lr = 1e-5
history = fit(epochs, lr, model, train_loader, val_loader)

I can see that the model is good and that the loss decreases. The problem is when I have to do a prediction on an example:
def predict_single(input, target, model):
    inputs = input.unsqueeze(0)
    predictions = model(inputs)
    prediction = predictions[0].detach()
    print(&quot;Input:&quot;, input)
    print(&quot;Target:&quot;, target)
    print(&quot;Prediction:&quot;, prediction)
    return prediction

input, target = val_df[1]
prediction = predict_single(input, target, model)

This returns:
Input: tensor([0.8705, 0.3900, 2.1000, 0.0650, 4.1206, 3.3000, 0.5300, 0.2610])
Target: tensor([6.])
Prediction: tensor([ 3.6465,  0.2800, -0.4561, -1.6733, -0.6519, -0.1650])

I want to see what are associated these logits, in the sense that I know that the highest logit is associated with the predicted class, but I want to see that class. I also applied softmax to rescale these values in a probability:
prediction = F.softmax(prediction)
print(prediction)
output = model(input.unsqueeze(0))
_,pred = output.max(1)
print(pred)

And the output is the following:
tensor([0.3296, 0.1361, 0.1339, 0.1324, 0.1335, 0.1346])
tensor([0])

I don't know what is that tensor([0]). I expect my predicted label, a value like 6.1 if the target is 6. But I am not able to obtain this.

First of all, Lets review the way you are calculating loss. From your code:
loss = loss_fn(out,torch.argmax(targets, dim=1))

you are using torch.argmax function which expects targets size as torch.Size([num_samples, num_classes]) or torch.Size([32, 6]) in your case. Are you sure your training labels are compatible with this size? From your writing I understand that you are reading label class as a number (from 3 to 8). So, its size is torch.Size([32, 1]). So, when you are calling torch.argmax with training data, `torch.argmax' is always returning 0.
That's why the model is learning to predict the class 0 whatever the input is.
Now, as your class label (for training) is from 3 to 8. Unfortunately if we use these labels with your loss_fn or torch.nn.CrossEntropyLoss(), it will be matched with total 9 labels, (class0 to class8) as maximum class labels is 8. So, you need to transform 3 to 8 -&gt; 0 to 5. For loss calculation use:
loss = loss_fn(out, targets - 3)

