A Beginner’s Guide to PyTorch: Building Your First Neural Network
Introduction
Welcome to the world of deep learning! If you’re just starting, PyTorch is one of the best frameworks to begin your journey. Known for its flexibility and ease of use, PyTorch has become a favorite among researchers and developers in the industry. In this blog, I’ll walk you through building your first neural network using PyTorch, from setting up your environment to training and evaluating your model.
What is PyTorch?
PyTorch is an open-source deep learning framework developed by Facebook’s AI Research lab. It’s known for its dynamic computation graph, which allows for flexible model building and debugging. PyTorch’s user-friendly interface and extensive community support make it an excellent choice for both beginners and experts.
Setting Up Your Environment
Before we dive into coding, let’s set up our environment.
Installation
First, you need to install PyTorch and torchvision, a package containing popular datasets, model architectures, and image transformations for computer vision.
pip install torch torchvision# verify your installation by -
import torch
print(torch.__version__)Setting Up a Virtual Environment (Optional)
Using a virtual environment can help you manage dependencies and avoid conflicts with other projects.
python -m venv pytorch-env
source pytorch-env/bin/activate # On Windows use `pytorch-env\Scripts\activate`- python -m venv pytorch-env: Creates a new virtual environment named
pytorch-env. - source pytorch-env/bin/activate: Activates the virtual environment. On Windows, use
pytorch-env\Scripts\activate.
Understanding Tensors
Tensors are the fundamental building blocks of PyTorch. They are similar to NumPy arrays but come with added functionality for GPU acceleration.
Creating a Tensor
import torch
# Creating a tensor
x = torch.tensor([1, 2, 3]) # Creates a 1D tensor with the values 1, 2,and 3.
print(x)Basic Tensor Operations
y = torch.tensor([4, 5, 6])
print(x + y)Building Your First Neural Network
Let’s build a simple neural network to classify handwritten digits from the MNIST dataset.
Data Preparation
We need to load and preprocess the MNIST dataset. In PyTorch, this can be done efficiently using the torchvision library, which provides utilities for various transformations and dataset loading.
import torchvision
import torchvision.transforms as transforms
# Define a transformation pipeline to convert images to tensors and normalize them
transform = transforms.Compose([
transforms.ToTensor(), # Converts a PIL image or NumPy array to a PyTorch tensor
transforms.Normalize((0.5,), (0.5,)) # Normalizes tensor image with mean 0.5 and std 0.5
])
# Load the MNIST dataset with the defined transformations
trainset = torchvision.datasets.MNIST(
root='./data', # Directory to store the dataset
train=True, # Load the training set
download=True, # Download the dataset if not already present
transform=transform # Apply the transformations to the dataset
)
# Create a DataLoader to iterate over the dataset
trainloader = torch.utils.data.DataLoader(
trainset, # The dataset to load
batch_size=32, # Number of samples per batch
shuffle=True # Shuffle the data at every epoch
)To make this more concrete, let’s look at a small example using these transforms:
from PIL import Image
import numpy as np
# Create a dummy image (5x5 pixels, grayscale)
dummy_image = Image.fromarray(np.random.randint(0, 256, (5, 5), dtype=np.uint8))
# Apply the transforms
tensor_image = transforms.ToTensor()(dummy_image)
normalized_image = transforms.Normalize((0.5,), (0.5,))(tensor_image)
print(f"Original Image: \n{np.array(dummy_image)}")
print(f"Tensor Image: \n{tensor_image}")
print(f"Normalized Image: \n{normalized_image}")Original Image:
[[109 228 167 50 120]
[226 48 158 60 85]
[ 5 191 183 83 105]
[142 212 121 104 171]
[ 35 42 114 247 88]]
Tensor Image:
tensor([[[0.4275, 0.8941, 0.6549, 0.1961, 0.4706],
[0.8863, 0.1882, 0.6196, 0.2353, 0.3333],
[0.0196, 0.7490, 0.7176, 0.3255, 0.4118],
[0.5569, 0.8314, 0.4745, 0.4078, 0.6706],
[0.1373, 0.1647, 0.4471, 0.9686, 0.3451]]])
Normalized Image:
tensor([[[ -0.1451, 0.7882, 0.3098, -0.6078, -0.0588],
[ 0.7725, -0.6235, 0.2392, -0.5294, -0.3333],
[ -0.9608, 0.4980, 0.4353, -0.3490, -0.1765],
[ 0.1137, 0.6627, -0.0510, -0.1843, 0.3412],
[ -0.7255, -0.6706, -0.1059, 0.9373, -0.3098]]])By chaining these transformations, we ensure that our data is in the right format and properly scaled for training the neural network.
Building the model
Let’s define a simple neural network.
import torch.nn as nn
import torch.nn.functional as F
# Define the model
class firstNN(nn.Module):
def __init__(self):
super(firstNN, self).__init__()
self.fc1 = nn.Linear(28 * 28, 128) # First fully connected layer with input size 28x28 and output size 128
self.fc2 = nn.Linear(128, 10) # Second fully connected layer with input size 128 and output size 10 (number of classes)
def forward(self, x):
x = x.view(-1, 28 * 28) # Flatten the input tensor
x = F.relu(self.fc1(x)) # Apply ReLU activation to the output of the first layer
x = self.fc2(x) # Compute the output of the second layer
return x
# Instantiate the model
model = firstNN()
print(model)SimpleNN(
(fc1): Linear(in_features=784, out_features=128, bias=True)
(fc2): Linear(in_features=128, out_features=10, bias=True)
)Training the model
Let’s define the loss function and optimizer, and then train the model.
import torch.optim as optim
# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss() # Loss function for multi-class classification
optimizer = optim.SGD(model.parameters(), lr=0.01) # Stochastic Gradient Descent (SGD) with learning rate 0.01
# Train the model for 5 epochs
for epoch in range(5): # Loop over the dataset multiple times
running_loss = 0.0
for inputs, labels in trainloader:
optimizer.zero_grad() # Clear the gradients of all optimized tensors
outputs = model(inputs) # Compute the forward pass
loss = criterion(outputs, labels) # Compute the loss
loss.backward() # Compute the gradient of the loss with respect to model parameters
optimizer.step() # Update the model parameters based on the gradients
running_loss += loss.item() # Accumulate the loss for the epoch
print(f'Epoch {epoch+1}, Loss: {running_loss/len(trainloader)}') # Print the average loss for the epochEpoch 1, Loss: 0.5623
Epoch 2, Loss: 0.3421
Epoch 3, Loss: 0.2754
Epoch 4, Loss: 0.2345
Epoch 5, Loss: 0.2046Evaluating the model
Finally, let’s evaluate the model’s performance.
correct = 0
total = 0
with torch.no_grad(): # Disable gradient computation for evaluation
for inputs, labels in trainloader:
outputs = model(inputs) # Compute the forward pass
_, predicted = torch.max(outputs.data, 1) # Find the predicted class with the highest score
total += labels.size(0) # Increment the total sample count
correct += (predicted == labels).sum().item() # Increment the correct prediction count
print(f'Accuracy: {100 * correct / total}%') # Print the model accuracyAccuracy: 93.4%Visualizing our results
import matplotlib.pyplot as plt
epochs = range(1, 6)
losses = [0.5623, 0.3421, 0.2754, 0.2345, 0.2046]
# Plot the loss values against epochs
plt.plot(epochs, losses, label='Training loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.xticks(epochs)
plt.legend()
plt.show() 
This plot visually represents how the training loss decreases as the number of epochs increases, indicating that the model is learning.
Congratulations! You’ve built and trained your first neural network using PyTorch. We’ve covered the basics of setting up your environment, understanding tensors, building a simple neural network, training the model, evaluating its performance, and visualizing the training progress. This is just the beginning of your deep learning journey with PyTorch. As you progress, you can explore more complex architectures, optimization techniques, and real-world applications.
Cheers!
Connect with me on Linkedin
View my work here
Visit my portfolio
