Beyond Raw Scores: Unveiling the Power of Predicted Probabilities in PyTorch

In classification tasks using PyTorch models, the model often outputs raw scores (logits) for each possible class. These scores represent the model's preference for each class, but they're not directly interpretable as probabilities.

The Softmax Function

To convert logits into probabilities, we employ the softmax function. Softmax takes a vector of logits as input and transforms it into a vector of probabilities between 0 and 1, where the sum of all probabilities equals 1. Each element in the output vector signifies the probability of the corresponding class.

2. Apply Softmax: Use the nn.functional.softmax function from PyTorch's functional module. Here's the syntax:

   import torch
   from torch import nn
   
   # Assuming your model's output is stored in 'logits'
   probabilities = nn.functional.softmax(logits, dim=-1)
   

   • logits: The tensor containing the model's raw predictions (logits).
   • dim=-1: The dimension along which the softmax operation should be applied. Here, we apply it across the last dimension (typically the class dimension) to ensure probabilities for each class.

Example Code

import torch
from torch import nn
import torch.nn.functional as F  # Import for softmax

# Sample model (replace with your actual model)
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # ... your model architecture here ...

    def forward(self, x):
        # ... your model's forward pass logic ...
        logits = x  # Assuming the last layer outputs logits
        return logits

# Create an instance of your model
model = MyModel()

# Sample input data (replace with your actual input)
input_data = torch.randn(1, 10)  # Batch size 1, feature vector of size 10

# Forward pass
logits = model(input_data)

# Get predicted probabilities
probabilities = F.softmax(logits, dim=-1)

print(probabilities)  # Output will be a tensor of probabilities for each class

Interpretation

The probabilities tensor will now hold the model's predicted probabilities for each class in your classification problem. You can use these probabilities to make more informed decisions, such as:

• Selecting the class with the highest probability as the predicted class.
• Setting a threshold for probability (e.g., only consider classes with a probability above 0.8).
• Visualizing the probabilities to gain insights into the model's confidence in its predictions.

import torch
from torch import nn
import torch.nn.functional as F

class BinaryClassifier(nn.Module):
    def __init__(self):
        super(BinaryClassifier, self).__init__()
        self.fc1 = nn.Linear(10, 1)  # Input size 10, output size 1 (logit)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return x  # Assuming the last layer outputs logits

# Create model and input
model = BinaryClassifier()
input_data = torch.randn(1, 10)  # Batch size 1, feature vector of size 10

# Forward pass
logits = model(input_data)

# Get predicted probabilities
probabilities = F.softmax(logits, dim=-1)

print(probabilities)  # Output will be a tensor with two probabilities (class 0 and 1)

Multi-class Classification (More than 2 Classes):

import torch
from torch import nn
import torch.nn.functional as F

class MultiClassClassifier(nn.Module):
    def __init__(self, num_classes):
        super(MultiClassClassifier, self).__init__()
        self.fc1 = nn.Linear(10, num_classes)  # Input size 10, output size = num_classes (logits)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return x  # Assuming the last layer outputs logits

# Create model with 5 classes and input
num_classes = 5
model = MultiClassClassifier(num_classes)
input_data = torch.randn(1, 10)  # Batch size 1, feature vector of size 10

# Forward pass
logits = model(input_data)

# Get predicted probabilities
probabilities = F.softmax(logits, dim=-1)

print(probabilities)  # Output will be a tensor with probabilities for all 5 classes

Using a Pre-trained Model (Example with ResNet):

import torch
from torch import nn
from torchvision import models

# Load a pre-trained ResNet model
model = models.resnet18(pretrained=True)

# Modify the last layer to output desired number of classes
num_classes = 10  # Adjust according to your classification task
model.fc = nn.Linear(model.fc.in_features, num_classes)

# Forward pass (assuming you have prepared your input data)
# ... your forward pass logic here ...

# Get predicted probabilities after the final layer
probabilities = F.softmax(logits, dim=-1)

print(probabilities)

• The sigmoid function, also known as the logistic function, can be used for binary classification tasks (two classes) as an alternative to softmax. It squashes values between 0 and 1, representing probabilities for each class. However, sigmoid is less numerically stable than softmax and might not be preferred in all cases.

import torch
from torch import nn

# Assuming your model's output is stored in 'logits'
probabilities = torch.sigmoid(logits)

Custom Output Layer (for Specific Probability Distributions):

• If your classification problem demands a specific probability distribution beyond a simple multinomial distribution (softmax output), you can create a custom output layer in your PyTorch model. This layer would implement the desired probability distribution's calculations to generate the probabilities directly. Here's an example structure:

class CustomDistributionOutput(nn.Module):
    def __init__(self, num_classes, distribution_type):
        super(CustomDistributionOutput, self).__init__()
        # ... define layers based on the chosen distribution_type ...

    def forward(self, x):
        # ... implement calculations for the chosen distribution ...
        probabilities = self.distribution(x)  # Replace with your distribution logic
        return probabilities

# Example usage:
model = MyModel(..., output_layer=CustomDistributionOutput(num_classes, "gamma"))

Temperature Scaling (Adjusting Softmax Output):

• Temperature scaling involves applying a temperature parameter (T) to the logits before applying softmax. This can be used to:
  • Increase model confidence (higher T): Softmax outputs become more concentrated on the most likely class.
  • Decrease model confidence (lower T): Softmax outputs become more spread out, considering multiple potential classes.

import torch
from torch import nn

temperature = 2.0  # Adjust as needed
probabilities = nn.functional.softmax(logits / temperature, dim=-1)

Choosing the Right Method:

• Softmax is the most general and widely used approach for obtaining predicted probabilities in PyTorch classification tasks.
• Sigmoid is a simpler alternative for binary classification but might have numerical stability issues.
• Custom output layers provide flexibility for specific probability distributions but require more development effort.
• Temperature scaling can fine-tune the confidence level of softmax outputs.