Building Your First Neural Network with PyTorch

Neural networks are the backbone of modern AI — from image classifiers to the transformers powering ChatGPT and Claude. In this tutorial, you will build, train, and evaluate a complete image classifier in PyTorch from scratch. No magic wrappers, no shortcuts: just the full pipeline so you understand every step.

By the end, you will have a working model that classifies handwritten digits from the MNIST dataset with over 97% accuracy, running on your CPU or GPU.

Prerequisites

• Basic Python knowledge (functions, loops, classes)
• Familiarity with NumPy-style arrays is helpful
• Install: pip install torch torchvision
• A GPU accelerates training but is not required for MNIST.

Why PyTorch?

PyTorch is the dominant framework in both academic research and production AI engineering. Its 'define-by-run' (eager execution) model means you write normal Python code — loops, conditionals, print statements — and it just works. This makes debugging dramatically easier than older static graph frameworks. All major AI labs publish their model code in PyTorch first.

1. Tensors: The Foundation

PyTorch's core data structure is the Tensor — an n-dimensional array identical to NumPy, but with two superpowers: GPU acceleration and automatic gradient tracking. Every computation in a neural network is a tensor operation.

import torch

# Create a tensor
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# Move to GPU if available — one line is all it takes
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)
print(f"Tensor is on {x.device}")

2. Defining the Model with nn.Module

In PyTorch, every neural network is a Python class that inherits from nn.Module. You define the layers in __init__ and the forward pass (how data flows through those layers) in forward. PyTorch automatically computes gradients through this graph during backpropagation.

For MNIST (28x28 grayscale images → 10 digit classes), we use a Multi-Layer Perceptron (MLP):

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

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()          # 28x28 image -> 784 vector
        self.fc1 = nn.Linear(28 * 28, 128)   # 784 inputs -> 128 hidden units
        self.fc2 = nn.Linear(128, 10)         # 128 -> 10 output classes

    def forward(self, x):
        x = self.flatten(x)
        x = F.relu(self.fc1(x))  # ReLU activation adds non-linearity
        x = self.fc2(x)           # Raw scores (logits) — no softmax yet
        return x

model = SimpleMLP().to(device)

The ReLU activation function is critical. Without it, stacking linear layers is mathematically equivalent to a single linear layer — the network cannot learn complex patterns.

3. Data Loading

torchvision provides pre-built datasets and transforms. The Normalize transform standardizes pixel values using the known MNIST mean (0.1307) and standard deviation (0.3081), which speeds up training significantly.

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

The DataLoader handles batching and shuffling automatically. A batch size of 64 is a solid default for MNIST.

4. The Training Loop

This is the heart of PyTorch. Every batch follows the same four-step rhythm: zero gradients, forward pass, compute loss, backward pass.

import torch.optim as optim

criterion = nn.CrossEntropyLoss()  # Combines softmax + negative log-likelihood
optimizer = optim.Adam(model.parameters(), lr=0.001)

model.train()
for epoch in range(3):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()                    # 1. Clear old gradients
        output = model(data)                     # 2. Forward pass
        loss = criterion(output, target)         # 3. Compute loss
        loss.backward()                          # 4. Backward pass (autograd)
        optimizer.step()                         # 5. Update weights
        
        if batch_idx % 200 == 0:
            print(f"Epoch {epoch} | Batch {batch_idx} | Loss: {loss.item():.4f}")

Why zero_grad()? PyTorch accumulates gradients by default. If you skip this step, gradients from the previous batch are added to the current batch's gradients, corrupting your updates.

5. Evaluation

Always evaluate with model.eval() (disables dropout and batchnorm randomness) and torch.no_grad() (skips gradient tracking to save memory and compute).

test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=1000)

model.eval()
correct = 0
with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()

print(f"Test Accuracy: {100. * correct / len(test_dataset):.1f}%")

You should see accuracy above 97% after just 3 epochs. From here, try replacing the MLP with a CNN (using nn.Conv2d) — it will push accuracy above 99% on the same dataset.