Skip to content

Cross-Entropy Loss Function in PyTorch

Using the Cross-Entropy Loss Function in PyTorch Cover Image

In this tutorial, you’ll learn about the Cross-Entropy Loss Function in PyTorch for developing your deep-learning models. The cross-entropy loss function is an important criterion for evaluating multi-class classification models.

This tutorial demystifies the cross-entropy loss function, by providing a comprehensive overview of its significance and implementation in deep learning. By the end of this tutorial, you’ll have learned the following:

  • How the cross-entropy loss function is calculated,
  • How the cross-entropy loss function is implemented in PyTorch,
  • When to use the cross-entropy loss function in your deep-learning models

If you’re looking for a quick answer, cross-entropy loss can be implemented in PyTorch for developing deep learning models as shown below:

# Implementing Cross-Entropy Loss in PyTorch
import torch
import torch.nn as nn

criterion = nn.CrossEntropyLoss()
loss = criterion(preds, labels)

Understanding the Cross-Entropy Loss Function

Loss functions are essential for guiding model training and enhancing the predictive accuracy of models. The cross-entropy loss function is a fundamental concept in classification tasks, especially in multi-class classification. The tool allows you to quantify the difference between predicted probabilities and the actual class labels.

Entropy is based on information theory, measuring the amount of uncertainty or randomness in a given probability distribution. You can think of it as measuring how uncertain we are about the outcomes of a random variable, where high entropy indicates more randomness while low entropy indicates more predictability.

Cross-entropy is an extension of entropy that allows you to quantify the difference between two probability distributions. In classification tasks, for example, one distribution represents the predicted probabilities assigned by a model to various classes, while the other distribution represents the true class labels. Cross-entropy, then, measures how similar the predicted probabilities are to the actual labels, by providing a numerical measure of dissimilarity.

The formula for cross-entropy loss is:

Cross-Entropy Loss = -∑(yᵢ * log(pᵢ))

Where the variables are defined as below:

  • y is the true class label (0 or 1 for binary classification, one-hot encoded vector for multi-class classification).
  • p is the predicted probability of the positive class in binary classification or the predicted probability vector in multi-class classification.
  • yi​ is the true probability of the i-th class (equal to 1 for the true class, 0 otherwise) in multi-class classification.
  • N is the number of classes in multi-class classification.
  • log denotes the natural logarithm.

Let’s now take a look at how the cross-entropy loss function is implemented in PyTorch.

Implementing Cross-Entropy Loss in PyTorch

In this section, we’ll bridge the gap between theory and practice by demonstrating the hands-on implementation of cross-entropy loss using PyTorch. We’ll cover the core concepts required to construct a classification model, compute predicted probabilities (logits), and calculate the cross-entropy loss.

In PyTorch, the cross-entropy loss function is implemented using the nn.CrossEntropyLoss class. Let’s take a look at how the class can be implemented. We’ll start by defining two variables: one containing sample predictions along multiple classes and another containing our true labels.

# Import Libraries
import torch
import torch.nn as nn

# Define 3 Samples and 4 Classes
preds = torch.randn(3, 4)

# Define Target Class Indices
labels = torch.LongTensor([0, 1, 2])

print(f'Input data: \n{preds}')

# Returns:
# Input data: 
# tensor([[ 0.3703, -0.1607, -0.7473, -2.2135],
#         [ 0.6929, -0.2042, -0.7660, -0.3613],
#         [-0.3363,  0.6417,  0.3103,  0.3132]])
# Labels:
# tensor([0, 1, 2])

Let’s break down what we did in the code block above:

  1. We imported our dependencies, namely torch itself and the nn module
  2. We then created our predictions tensor, which contains 3 samples, each with four outputs. This replicates the predicted value for each of the four classes
  3. We also created a LongTensor representing the different classes. These are made up, but represent what our classes could be. If you want to modify the data, just don’t have it exceed the number of classes there are (in this case 4, bringing us to a max value of 3).

Now that we have our mock data defined, let’s see how we can implement cross-entropy loss to evaluate the loss in our predictions:

# Define Loss Function
criterion = nn.CrossEntropyLoss()
loss = criterion(preds, labels)

print(f'Loss: {loss}')

# Returns:
# Loss: 1.2129982709884644

In the code block above, we instantiated the CrossEntropyLoss class, allowing us to use this function to accept both prediction and label tensors. We then calculate the loss by passing in our predictions and labels. By printing the loss, we can see that the loss is around 1.21.

Interpreting the Cross-Entropy Loss for Deep Learning

Interpreting the loss value, particularly the cross-entropy loss, requires an understanding of its magnitude, trends, and context within your specific deep-learning problem. Here are some key points to consider when interpreting the loss value:

  1. Magnitude: A lower loss value indicates that predictions are close to their true, expected values. Inversely, a larger loss value means that the predictions are less accurate.
  2. Understanding baseline: Consider what a reasonable baseline loss might be for your problem. For example, random guessing in a binary classification problem might result in a loss around -log(0.5) ≈ 0.693. Your model’s loss should ideally be significantly lower than this baseline.
  3. Training vs. Validation Loss: if the loss on your training dataset is much lower than on your validation dataset, this could indicate that your model is overfitting.

Let’s now take a look at when to use the cross-entropy loss in deep learning projects.

When to Use Cross-Entropy Loss in Deep Learning

Cross-entropy loss is a great choice when working with classification problems. This is true for both binary classification problems (2 classes) or multi-class classification problems. It effectively measures the difference between predicted probabilities and true class labels, aligning with the goal of accurate classification. It integrates very well with the softmax activation function, allowing you to find non-linear patterns in your data.

The loss function also allows you to have a clear, understandable and continuous measure to understand the performance of your model. This allows you to better understand the model’s performance over training.

It also addresses the vanishing gradient problem more effectively than other cost functions like mean squared error, which can greatly blow errors out.

Finally, in cases where classes are imbalanced (one class has significantly fewer samples than others), cross-entropy loss can help mitigate the impact of class imbalance. It assigns higher weights to misclassifications of the minority class, leading to improved performance. This can be very helpful when you are collecting data for transfer learning and have smaller datasets.


In this tutorial, you’ve gained a solid understanding of the Cross-Entropy Loss Function in PyTorch, a fundamental tool for developing accurate deep-learning models, particularly in classification tasks. The journey through this tutorial has demystified the cross-entropy loss function, taking you from its conceptual significance to hands-on implementation in PyTorch.

Here’s a recap of what you’ve learned:

  • You explored the motivation behind cross-entropy loss and its importance in evaluating multi-class classification models.
  • You learned how to calculate and implement the cross-entropy loss function using PyTorch’s powerful capabilities.
  • You discovered the scenarios where cross-entropy loss is the ideal choice for your deep-learning models.

By mastering the concepts covered in this tutorial, you’re now equipped with the knowledge and tools to confidently apply cross-entropy loss in your deep learning projects. You can use this loss function to guide your model’s training, improve its predictive accuracy, and make informed decisions about model architecture and optimization strategies.

Remember that while cross-entropy loss is a versatile and effective tool, it’s essential to consider the unique characteristics of your dataset and problem domain. Keep refining your understanding of loss functions and their role in model development, and you’ll be well on your way to building high-performing deep-learning models. Happy coding and exploring the world of deep learning!

To learn more about cross-entropy loss in PyTorch, check out the official documentation.

Nik Piepenbreier

Nik is the author of and has over a decade of experience working with data analytics, data science, and Python. He specializes in teaching developers how to use Python for data science using hands-on tutorials.View Author posts

Leave a Reply

Your email address will not be published. Required fields are marked *