PyTorch Tutorial: Train Your First Neural Network

Learn how to build and train your first deep learning neural network from scratch using Python and PyTorch with step-by-step code examples.

Deep learning has revolutionized artificial intelligence, powering everything from computer vision systems to advanced large language models. At the heart of this technological shift lies PyTorch, an open-source machine learning framework developed by Meta that provides maximum flexibility, imperative execution, and seamless GPU acceleration.

This comprehensive tutorial will guide you through the process of building, training, and evaluating your very first feedforward neural network in Python using PyTorch.

1. Prerequisites and Environment Setup

Before writing any neural network code, ensure you have Python installed and set up a clean virtual environment. Install PyTorch along with torchvision for dataset utilities.

# Create and activate a virtual environment
python -m venv venv
source venv/bin/activate  # On Windows use: venv\Scripts\activate

# Install PyTorch and torchvision (CPU version)
pip install torch torchvision

2. Understanding PyTorch Tensors

At the core of PyTorch is the **Tensor**, a specialized data structure very similar to NumPy arrays, with the added capability of running on GPUs for accelerated hardware computation.

import torch

# Create a 2D tensor (matrix)
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print("Tensor x:\n", x)
print("Tensor Shape:", x.shape)
print("Tensor Device:", x.device)

3. Defining the Neural Network Architecture

In PyTorch, custom neural networks are built by subclassing `torch.nn.Module`. You define your layers in the `__init__` constructor and specify the forward pass logic in the `forward` method.

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

class SimpleNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Apply ReLU activation function to hidden layer
        x = F.relu(self.fc1(x))
        # Output layer
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleNN(input_dim=4, hidden_dim=8, output_dim=2)
print(model)

4. Writing the Training Loop

Training a neural network involves four main steps during each epoch: forward propagation, loss calculation, backward propagation (calculating gradients), and optimizer weight updates.

import torch.optim as optim

# Dummy training data (features and binary labels)
inputs = torch.randn(100, 4)
targets = torch.randint(0, 2, (100,))

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training loop epochs
epochs = 5
for epoch in range(epochs):
    # 1. Zero the gradients
    optimizer.zero_grad()
    
    # 2. Forward pass
    outputs = model(inputs)
    
    # 3. Calculate loss
    loss = criterion(outputs, targets)
    
    # 4. Backward pass and optimize
    loss.backward()
    optimizer.step()
    
    print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

Conclusion

You have successfully built, configured, and trained your first neural network in PyTorch! From handling tensors to designing custom `nn.Module` architectures and executing training loops, these foundational skills pave the way for tackling advanced deep learning challenges.