Train a Deep Learning Model With Pytorch
Last Updated :
23 Jul, 2025
Neural Network is a type of machine learning model inspired by the structure and function of human brain. It consists of layers of interconnected nodes called neurons which process and transmit information. Neural networks are particularly well-suited for tasks such as image and speech recognition, natural language processing and making predictions based on large amounts of data.
Below is an image of a simple neural network:
Neural NetworkIn this article we will train a neural network on the MNIST dataset. It is a dataset of handwritten digits consisting of 60,000 training examples and 10,000 test examples. Each example is a 28x28 grayscale image of a handwritten digit with values ranging from 0 (white) to 255 (black). The label for each example is the digit that the image represents with values ranging from 0 to 9. It is a dataset commonly used for training and evaluating image classification models particularly in the field of computer vision.
An sample data from MnistInstalling PyTorch
To install PyTorch you will need to have Python and pip (the package manager for Python) installed on your computer. We can install PyTorch and necessary libraries by running the following command in your command prompt or terminal:
pip install torch
pip install torchvision
pip install torchsummary
Implementing Deep Neural Network
The model will be a simple feedforward neural network and we’ll walk through the steps to implement it.
1. Importing necessary Libraries
We will be importing Pytorch and matplotlib.
- torch: PyTorch package for tensor operations.
- torchvision: A package for vision related datasets and models.
Python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
2. Loading MNIST Datasets
First we need to import the necessary libraries and load the dataset which can be loaded using the torchvision library. We will then convert images to tensors and normalizes their pixel values to have a mean of 0.5 and standard deviation of 0.5.
- transforms.Compose: A sequence of transformations to apply to the images.
- ToTensor(): Converts images into PyTorch tensors.
- Normalize(): Normalizes the pixel values to the range [-1, 1].
- DataLoader: DataLoader is a utility class that handles the batching, shuffling and loading of data for training and evaluation.
- batch_size=64 means that in each iteration 64 samples will be processed at once.
Python
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
3. Building model
Let’s define a simple feedforward neural network with one hidden layer. This model will take 28x28 images, flatten them and pass them through a fully connected layer followed by an output layer with 10 units (for the 10 digits).
- nn.Module: The base class for all neural network modules in PyTorch. By inheriting from this class we create a custom model with layers and a forward pass.
- nn.Linear: This is a basic layer where each input is connected to every output node.
- fc1: The first fully connected layer transforms the 28x28 image (flattened to a 784-length vector) into a 128-dimensional vector.
- fc2: The second fully connected layer outputs a 10-dimensional vector corresponding to the 10 possible digit classes.
- We use ReLU activation after the first layer to introduce non-linearity.
Python
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(-1, 28*28)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleNN()
4. Defining the Loss Function and Optimizer
For classification tasks we use CrossEntropyLoss as the loss function which is appropriate for multi-class classification. The Adam optimizer will be used to update the weights during training.
- CrossEntropyLoss: Measures the difference between the predicted and true labels.
- Adam: A popular optimization algorithm that adapts the learning rate during training.
- lr=0.001 sets the learning rate which determines how much we adjust the weights during each update.
Python
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
5. Training the Model
Now we will train the model. In each epoch we will perform a forward pass to compute predictions, calculate the loss then perform backpropagation to compute gradients. Finally we update the model weights using the optimizer. After every 100 mini-batches the average loss for that mini-batch is printed. The total loss for each epoch is also tracked for later visualization.
- loss.backward(): Computes gradients for each model parameter.
- optimizer.step(): Updates the model's weights using the gradients.
- cuda(): If a GPU is available we move the images, labels and model to the GPU for faster computation.
- optimizer.zero_grad(): Before performing a backward pass we need to set all the gradients to zero. Otherwise the gradients will accumulate from previous iterations.
- loss.backward(): This computes the gradients for the model’s weights indicating how much the weights need to be adjusted.
- optimizer.step(): This updates the model’s weights based on the gradients calculated during the backward pass.
Python
train_losses = []
num_epochs = 5
for epoch in range(num_epochs):
running_loss = 0.0
for i, (images, labels) in enumerate(trainloader, 0):
if torch.cuda.is_available():
images, labels = images.cuda(), labels.cuda()
model.cuda()
else:
model.cpu()
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 200 == 99:
print(f"[Epoch {epoch+1}, Batch {i+1}] Loss: {running_loss / 100:.4f}")
running_loss = 0.0
train_losses.append(running_loss / len(trainloader))
Output:
Train the Model6. Ploting Training and Validation Curves
The Training Loss Curve shows how the model's loss decreases over epochs indicating that the model is learning and improving its performance. It is useful to visualize how the training loss changes over time. This helps us understand if the model is learning effectively.
Python
plt.plot(train_losses, label="Training Loss")
plt.title('Training Loss Curve')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()
Output:
Plot Training and Validation CurvesThe training loss decreases steadily across epochs with a sharper drop between epochs 1 and 2. A slight increase at epoch 3 suggests a temporary plateau but the loss resumes its decline by epoch 4 indicating overall improvement in model performance.
7. Evaluating the Model
Once the model is trained we evaluate its performance on the test dataset. We compute the accuracy by comparing the predicted labels with the true labels.
- model.eval(): Sets the model to evaluation mode disabling dropout layers.
- torch.no_grad(): Ensures that no gradients are calculated during evaluation saving memory.
Python
correct = 0
total = 0
model.eval()
with torch.no_grad():
for images, labels in testloader:
if torch.cuda.is_available():
images, labels = images.cuda(), labels.cuda()
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")
Output:
Test Accuracy: 97.48%
We achieved a good accuracy of 97.48% which is great for a deep learning model.
Similar Reads
Deep Learning Tutorial Deep Learning is a subset of Artificial Intelligence (AI) that helps machines to learn from large datasets using multi-layered neural networks. It automatically finds patterns and makes predictions and eliminates the need for manual feature extraction. Deep Learning tutorial covers the basics to adv
5 min read
Deep Learning Basics
Introduction to Deep LearningDeep Learning is transforming the way machines understand, learn and interact with complex data. Deep learning mimics neural networks of the human brain, it enables computers to autonomously uncover patterns and make informed decisions from vast amounts of unstructured data. How Deep Learning Works?
7 min read
Artificial intelligence vs Machine Learning vs Deep LearningNowadays many misconceptions are there related to the words machine learning, deep learning, and artificial intelligence (AI), most people think all these things are the same whenever they hear the word AI, they directly relate that word to machine learning or vice versa, well yes, these things are
4 min read
Deep Learning Examples: Practical Applications in Real LifeDeep learning is a branch of artificial intelligence (AI) that uses algorithms inspired by how the human brain works. It helps computers learn from large amounts of data and make smart decisions. Deep learning is behind many technologies we use every day like voice assistants and medical tools.This
3 min read
Challenges in Deep LearningDeep learning, a branch of artificial intelligence, uses neural networks to analyze and learn from large datasets. It powers advancements in image recognition, natural language processing, and autonomous systems. Despite its impressive capabilities, deep learning is not without its challenges. It in
7 min read
Why Deep Learning is ImportantDeep learning has emerged as one of the most transformative technologies of our time, revolutionizing numerous fields from computer vision to natural language processing. Its significance extends far beyond just improving predictive accuracy; it has reshaped entire industries and opened up new possi
5 min read
Neural Networks Basics
What is a Neural Network?Neural networks are machine learning models that mimic the complex functions of the human brain. These models consist of interconnected nodes or neurons that process data, learn patterns and enable tasks such as pattern recognition and decision-making.In this article, we will explore the fundamental
12 min read
Types of Neural NetworksNeural networks are computational models that mimic the way biological neural networks in the human brain process information. They consist of layers of neurons that transform the input data into meaningful outputs through a series of mathematical operations. In this article, we are going to explore
7 min read
Layers in Artificial Neural Networks (ANN)In Artificial Neural Networks (ANNs), data flows from the input layer to the output layer through one or more hidden layers. Each layer consists of neurons that receive input, process it, and pass the output to the next layer. The layers work together to extract features, transform data, and make pr
4 min read
Activation functions in Neural NetworksWhile building a neural network, one key decision is selecting the Activation Function for both the hidden layer and the output layer. It is a mathematical function applied to the output of a neuron. It introduces non-linearity into the model, allowing the network to learn and represent complex patt
8 min read
Feedforward Neural NetworkFeedforward Neural Network (FNN) is a type of artificial neural network in which information flows in a single direction i.e from the input layer through hidden layers to the output layer without loops or feedback. It is mainly used for pattern recognition tasks like image and speech classification.
6 min read
Backpropagation in Neural NetworkBack Propagation is also known as "Backward Propagation of Errors" is a method used to train neural network . Its goal is to reduce the difference between the modelâs predicted output and the actual output by adjusting the weights and biases in the network.It works iteratively to adjust weights and
9 min read
Deep Learning Models
Deep Learning Frameworks
TensorFlow TutorialTensorFlow is an open-source machine-learning framework developed by Google. It is written in Python, making it accessible and easy to understand. It is designed to build and train machine learning (ML) and deep learning models. It is highly scalable for both research and production.It supports CPUs
2 min read
Keras TutorialKeras high-level neural networks APIs that provide easy and efficient design and training of deep learning models. It is built on top of powerful frameworks like TensorFlow, making it both highly flexible and accessible. Keras has a simple and user-friendly interface, making it ideal for both beginn
3 min read
PyTorch TutorialPyTorch is an open-source deep learning framework designed to simplify the process of building neural networks and machine learning models. With its dynamic computation graph, PyTorch allows developers to modify the networkâs behavior in real-time, making it an excellent choice for both beginners an
7 min read
Caffe : Deep Learning FrameworkCaffe (Convolutional Architecture for Fast Feature Embedding) is an open-source deep learning framework developed by the Berkeley Vision and Learning Center (BVLC) to assist developers in creating, training, testing, and deploying deep neural networks. It provides a valuable medium for enhancing com
8 min read
Apache MXNet: The Scalable and Flexible Deep Learning FrameworkIn the ever-evolving landscape of artificial intelligence and deep learning, selecting the right framework for building and deploying models is crucial for performance, scalability, and ease of development. Apache MXNet, an open-source deep learning framework, stands out by offering flexibility, sca
6 min read
Theano in PythonTheano is a Python library that allows us to evaluate mathematical operations including multi-dimensional arrays efficiently. It is mostly used in building Deep Learning Projects. Theano works way faster on the Graphics Processing Unit (GPU) rather than on the CPU. This article will help you to unde
4 min read
Model Evaluation
Deep Learning Projects