Implementation of dropout in PyTorch

Dropout is a regularization technique commonly used in neural networks to prevent overfitting. Overfitting occurs when a neural network learns to perform exceptionally well on the training data but fails to generalize effectively to new, unseen data. Dropout helps address this issue by randomly dropping out (i.e., setting to zero) a fraction of neurons during each training iteration. This prevents any single neuron from becoming overly reliant on specific features, leading to a more robust and generalizable model.

During training, dropout is applied to individual neurons with a certain probability, usually denoted as p. This probability determines the likelihood of a neuron being dropped out. As a result, the network becomes less sensitive to the precise configuration of neurons and learns to rely on more diverse features. Its range is in between $0$ and $1$ .

In the animation above, we describe how dropout affects our model. We can see that after applying dropout some of the neurons, i.e., second and third neuron of second layer and first neuron of third layer, are dropped out. This helps prevent the network from relying too heavily on any individual neuron and encourages the learning of more generalized features.

Syntax

The syntax to add dropout in the neural network is as follows:

nn.Dropout(p=dropout_prob)

Implementation

Let’s explore how dropout is integrated into a neural network implementation using Pytorch, and how we can use dropout to improve the performance of the model.

Training

The provided code explains a neural network implementation with and without dropout for performance comparison. We import the key libraries, and a reproducible seed is set before generating a training dataset. We construct two neural networks — one without dropout and another with dropout. This distinction allows for evaluating the impact of dropout on model performance. Each network comprises two layers with ReLU activation, and the dropout model includes a dropout layer between them. The models are trained using stochastic gradient descent (SGD) optimization with a learning rate of 0.01. Training occurs in separate loops for both models, involving forward passes, loss computation via cross-entropy, backpropagation, and optimizer updates.

import torch
import torch.nn as nn
import torch.optim as optim
# Set a random seed for reproducibility
torch.manual_seed(41)
X = torch.rand(100, 10) * 2 - 1  # Generate data between -1 and 1
y = ((X[:, 0] + X[:, 1]) > 0).long()  # Classify based on sum of first two features
class DropoutDemoModel(nn.Module):
    def __init__(self, dropout_prob):
        super(DropoutDemoModel, self).__init__()
        self.dropout_prob = dropout_prob
        self.fc1 = nn.Linear(10, 64)
        self.dropout = nn.Dropout(p=self.dropout_prob)
        self.fc2 = nn.Linear(64, 2)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x
# Convert to PyTorch tensors
X_tensor = X.float()
y_tensor = y
# Define model with and without dropout
model_no_dropout = DropoutDemoModel(dropout_prob=0.0)
model_with_dropout = DropoutDemoModel(dropout_prob=0.5)
# Define loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer_no_dropout = optim.SGD(model_no_dropout.parameters(), lr=0.01)
optimizer_with_dropout = optim.SGD(model_with_dropout.parameters(), lr=0.01)
# Training loop without dropout
for epoch in range(100):
    optimizer_no_dropout.zero_grad()
    outputs = model_no_dropout(X_tensor)
    loss = criterion(outputs, y_tensor)
    loss.backward()
    optimizer_no_dropout.step()
# Training loop with dropout
for epoch in range(100):
    optimizer_with_dropout.zero_grad()
    outputs = model_with_dropout(X_tensor)
    loss = criterion(outputs, y_tensor)
    loss.backward()
    optimizer_with_dropout.step()
print("Training is completed")

Code explanation

Lines 11–17: We create a custom neural network class DropoutDemoModel that is inherited from nn.Module. It has two fully connected layers (nn.Linear), with hidden layers, sizes 10 -> 64 -> 2, representing the neural network’s architecture. A dropout layer (nn.Dropout) with a given dropout probability is added after the first hidden layer. In this example, the dropout probability is set to 0.5.
Lines 19–23: We define the forward method of class DropoutDemoModel. The input data x is passed through the first linear layer, followed by a ReLU activation function. The output of the activation is passed through the dropout layer. Finally, the output is passed through the second linear layer to produce the final logits.

Testing

We generate a test dataset using the same data distribution process, evaluate two neural network models on this dataset (one with dropout and one without), and print out the test loss and accuracy for both cases. By evaluating the models with and without dropout on the test data, the code provides insights into how dropout affects the model’s performance on unseen data.

# Generate a test dataset using the same data generation process
torch.manual_seed(42)
X_test = torch.rand(50, 10) * 2 - 1
y_test = ((X_test[:, 0] + X_test[:, 1]) > 0).long()
# Convert test data to PyTorch tensors
X_test_tensor = X_test.float()
y_test_tensor = y_test
# Evaluation function
def evaluate_model(model, X, y):
    model.eval()
    with torch.no_grad():
        outputs = model(X)
        loss = criterion(outputs, y)
        _, predicted = torch.max(outputs, 1)
        accuracy = (predicted == y).sum().item() / len(y)
    return loss.item(), accuracy
# Evaluate models on test data without dropout
no_dropout_loss, no_dropout_accuracy = evaluate_model(model_no_dropout, X_test_tensor, y_test_tensor)
print(f"No Dropout - Test Loss: {no_dropout_loss:.4f}, Accuracy: {no_dropout_accuracy:.4f}")
# Evaluate models on test data with dropout
dropout_loss, dropout_accuracy = evaluate_model(model_with_dropout, X_test_tensor, y_test_tensor)
print(f"With Dropout - Test Loss: {dropout_loss:.4f}, Accuracy: {dropout_accuracy:.4f}")

Code explanation

Lines 4–5: We generate the test dataset using the same distribution process.
Lines 12–19: We define a function evaluate_model to evaluate a given neural network model on the test data. Inside the function, the model is set to evaluation mode using model.eval() to disable dropout during evaluation. A forward pass is performed using the model to obtain predictions. The loss between the predictions and the true labels is computed using a predefined loss criterion. The predictions are compared to the true labels to calculate accuracy. Finally, the function returns the loss and accuracy as values.

Use cases

Here are some situations where it is recommended to use dropout:

Preventing overfitting: Dropout is primarily used to prevent overfitting.
Complex networks: Dropout is particularly useful when dealing with complex neural network architectures, such as deep neural networks with many layers.
Limited training data: When we have a limited amount of training data, overfitting becomes a significant concern. Dropout can help the model generalize better from the limited data by preventing it from memorizing the training examples.
High-dimensional data: If we’re working with high-dimensional data, where each input feature provides a lot of information, the risk of overfitting is higher. Dropout can help regularize the model and avoid capturing noise from the high-dimensional input space.

Conclusion

Free Resources

Learn in-demand tech skills in half the time

PRODUCTS

Mock Interview

New

Courses

Skill Paths

Projects

Assessments

	Model without Dropout	Model with Dropout
Test Loss	0.6574	0.6022
Accuracy	0.66	0.86

Implementation of dropout in PyTorch

Syntax

Implementation

Training

Code explanation

Testing

Model with Dropout vs. Model without Dropout

Code explanation

Use cases

Conclusion