03. Computer Vision - Why do we torch.argmax(dim=1) in logits->prediction probs->prediction labels?

[theanilsomani]

In the below code, I want to ask why do we use the torch.argmax(dim=1) at "Turn predictions from logits -> prediction probabilities -> predictions labels"
What is the reason behind it?
Make predictions with trained model y_preds = [] model_2.eval() with torch.inference_mode(): for X, y in tqdm(test_dataloader, desc="Making predictions"): # Send data and targets to target device X, y = X.to(device), y.to(device) # Do the forward pass y_logit = model_2(X) # Turn predictions from logits -> prediction probabilities -> predictions labels y_pred = torch.softmax(y_logit, dim=1).argmax(dim=1) # Put predictions on CPU for evaluation y_preds.append(y_pred.cpu())

[IIsameerII]

In a scenario where we don't specify the dim (dimension), pytorch would calculate the argmax for all the values as one.

In this case 0.8433 is the highest value so the index position is 10.

But if we specify the dim=1 then we are essentially telling pytorch to go into 1st dim of the tensor and get the argmax.

Also, we use torch.argmax(dim=1) in this case because the logits are all the 10 classes for each image in stored in dim 1.

[theanilsomani]

Correct me if I am wrong, you are saying that we use torch.argmax(dim=1) to find the maximum value across a single rows.

[Kabilduke]

For :
Text (dim = -1) to find maximum value in a sequence
Image (dim = 1) to find maximum value across a single rows