How to Build Your First Neural Network in Python

How to Build Your First Neural Network in Python

Building a neural network sounds intimidating, but modern frameworks like PyTorch make it surprisingly approachable. In this tutorial, you'll build a complete neural network from scratch that classifies handwritten digits — and you'll understand every line of code.

By the end, you'll have a working model that achieves over 97% accuracy on the classic MNIST dataset.

Prerequisites

You'll need Python 3.8+ and PyTorch installed:

pip install torch torchvision

Step 1: Load the Data

MNIST contains 70,000 grayscale images of handwritten digits (0–9), each 28×28 pixels. PyTorch's torchvision makes loading it trivial:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Transform images to tensors and normalize
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Download and load training + test data
train_data = datasets.MNIST("./data", train=True,
                            download=True, transform=transform)
test_data = datasets.MNIST("./data", train=False,
                           transform=transform)

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=1000)

What's happening: We normalize pixel values using MNIST's known mean (0.1307) and standard deviation (0.3081). The DataLoader handles batching and shuffling automatically.

Step 2: Define the Model

A neural network is a stack of layers that transform input data into predictions. We'll build a simple but effective feedforward network:

class NeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(28 * 28, 256),   # Input: 784 pixels
            nn.ReLU(),                  # Activation function
            nn.Dropout(0.2),            # Regularization
            nn.Linear(256, 128),        # Hidden layer
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 10)          # Output: 10 digit classes
        )

    def forward(self, x):
        x = self.flatten(x)            # 28x28 image → 784 vector
        return self.layers(x)

model = NeuralNet()
print(model)

Breaking it down:

• nn.Flatten() — Converts the 28×28 image into a flat vector of 784 values
• nn.Linear(in, out) — A fully connected layer that learns weights and biases
• nn.ReLU() — The Rectified Linear Unit activation: max(0, x). It introduces non-linearity, which is what allows neural networks to learn complex patterns
• nn.Dropout(0.2) — Randomly zeroes 20% of values during training to prevent overfitting
• The final layer outputs 10 values (one per digit class)

Step 3: Set Up Training

We need two things: a loss function to measure how wrong our predictions are, and an optimizer to update the model's weights:

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

CrossEntropyLoss is the standard choice for classification tasks. Adam is a reliable optimizer that adapts the learning rate for each parameter.

Step 4: Train the Model

Training means repeatedly showing the model batches of data, computing the loss, and updating weights:

def train(model, loader, criterion, optimizer, epochs=5):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        correct = 0
        total = 0

        for images, labels in loader:
            optimizer.zero_grad()           # Reset gradients
            outputs = model(images)         # Forward pass
            loss = criterion(outputs, labels)
            loss.backward()                 # Compute gradients
            optimizer.step()                # Update weights

            total_loss += loss.item()
            _, predicted = outputs.max(1)
            correct += predicted.eq(labels).sum().item()
            total += labels.size(0)

        acc = 100.0 * correct / total
        print(f"Epoch {epoch+1}: Loss={total_loss:.2f}, "
              f"Accuracy={acc:.1f}%")

train(model, train_loader, criterion, optimizer)

The training loop explained:

1. zero_grad() — Clear previous gradients (PyTorch accumulates them by default)
2. Forward pass — Run input through the model to get predictions
3. Compute loss — How far off are the predictions?
4. backward() — Compute gradients via backpropagation
5. step() — Update model weights using those gradients

Step 5: Evaluate on Test Data

After training, check how well the model generalizes to unseen data:

def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in loader:
            outputs = model(images)
            _, predicted = outputs.max(1)
            correct += predicted.eq(labels).sum().item()
            total += labels.size(0)

    print(f"Test Accuracy: {100.0 * correct / total:.1f}%")

evaluate(model, test_loader)

With this simple architecture and just 5 epochs of training, you should see ~97% test accuracy. Not bad for a first neural network!

Key Concepts Recap

| Concept | What It Does | |---------|-------------| | Forward pass | Input flows through layers to produce predictions | | Loss function | Measures prediction error (CrossEntropy for classification) | | Backpropagation | Computes how each weight contributed to the error | | Optimizer | Updates weights to reduce future errors | | Epochs | Number of complete passes through the training data |

Where to Go from Here

This feedforward network is a solid starting point. To improve further, explore:

• Convolutional Neural Networks (CNNs) — Purpose-built for image data, reaching 99%+ on MNIST
• Learning rate scheduling — Gradually reduce the learning rate during training
• Data augmentation — Artificially expand your training set with transformations
• GPU acceleration — Move your model and data to GPU with .to("cuda") for faster training

The patterns you've learned here — define a model, write a training loop, evaluate on test data — apply to virtually every deep learning project, from image classification to language models.

---

Want to go deeper? Check out our ML Basic Course →