Build Your First Neural Network With PyTorch, Step by Step

A neural network is a stack of differentiable functions whose weights are tuned by gradient descent. PyTorch lets you build one in about 40 lines of Python. This guide takes you from an empty file to a trained MNIST digit classifier with roughly 97–98% test accuracy, and explains every design choice along the way: tensors, layers, loss functions, optimizers, and the training loop.

Quick Takeaways

• PyTorch 2.x gives you dynamic computation graphs, autograd, and torch.compile with a Python-native API.
• A minimal network needs five parts: data pipeline, model (nn.Module), loss function, optimizer, training loop.
• A 784 → 128 → 64 → 10 MLP with ReLU and Adam (lr=1e-3) reaches ~97% on MNIST in 5 epochs on a CPU.
• Always call model.train() / model.eval() and optimizer.zero_grad(). Skipping them causes the most common beginner bugs.

What You Need Before You Start

You need Python 3.9+ and basic comfort with functions and classes. No GPU is required. MNIST trains in minutes on a laptop CPU.

Install PyTorch and torchvision:

pip install torch torchvision

For GPU builds, use the selector on pytorch.org to get the right CUDA command for your system.

Verify the install:

import torch
print(torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("MPS available:", torch.backends.mps.is_available())  # Apple Silicon

Step 1: Understand Tensors

A tensor is an n-dimensional array with two extra abilities: it runs on GPUs, and it tracks operations for automatic differentiation.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # shape (2, 2), dtype float32
w = torch.randn(2, 1, requires_grad=True)    # learnable weights

y = x @ w                                    # matrix multiply
loss = y.sum()
loss.backward()                              # computes d(loss)/dw

print(w.grad)  # gradients stored here

Three concepts matter here:

• requires_grad=True tells PyTorch to record operations on this tensor.
• .backward() runs backpropagation through the recorded graph.
• .grad holds the resulting gradient. PyTorch accumulates gradients across calls, which is why you must zero them each iteration.

The Core Math

Gradient descent updates each weight by a small step against its gradient:

w ← w − η · ∂L/∂w

Here η is the learning rate. Too high and training diverges. Too low and it crawls. 0.001 is a safe starting point for Adam.

Step 2: Load and Prepare the Data

MNIST contains 70,000 grayscale images at 28×28 pixels, labeled 0–9. We normalize pixel values and wrap the data in DataLoader objects for batching and shuffling.

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

transform = transforms.Compose([
    transforms.ToTensor(),                      # [0,255] -> [0.0,1.0], shape (1,28,28)
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean and std
])

train_data = datasets.MNIST("data", train=True, download=True, transform=transform)
test_data = datasets.MNIST("data", train=False, download=True, transform=transform)

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

Why normalize? Inputs centered near zero with unit variance keep gradients stable and speed up convergence. The values 0.1307 and 0.3081 are the dataset’s pixel mean and standard deviation.

Why shuffle only training data? Shuffling prevents the model from learning batch order. Evaluation order is irrelevant.

Step 3: Define the Model

Every PyTorch model subclasses nn.Module. You declare layers in __init__ and define the data flow in forward.

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()          # (B,1,28,28) -> (B,784)
        self.net = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),               # 10 raw scores (logits)
        )

    def forward(self, x):
        return self.net(self.flatten(x))

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
model = MLP().to(device)
print(model)

Architecture Breakdown

Each linear layer computes y = xWᵀ + b. ReLU (max(0, x)) adds the non-linearity. Without it, stacked linear layers collapse into a single linear transformation, and the network could never learn curved decision boundaries.

The final layer outputs logits, not probabilities. Do not add softmax here. CrossEntropyLoss applies log-softmax internally, which is more numerically stable.

Step 4: Choose a Loss Function and Optimizer

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

CrossEntropyLoss suits multi-class classification. It penalizes confident wrong answers heavily.

Adam adapts the learning rate per parameter using running averages of gradients and squared gradients. It is a strong default for first projects.

Step 5: Write the Training Loop

The loop repeats the same five operations for every batch.

def train_one_epoch(model, loader, loss_fn, optimizer):
    model.train()
    total_loss = 0.0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        logits = model(images)            # 1. forward pass
        loss = loss_fn(logits, labels)    # 2. compute loss

        optimizer.zero_grad()             # 3. clear old gradients
        loss.backward()                   # 4. backpropagate
        optimizer.step()                  # 5. update weights

        total_loss += loss.item() * images.size(0)
    return total_loss / len(loader.dataset)

Do not skip optimizer.zero_grad(). Gradients accumulate by default, so without it each step uses the sum of all previous gradients and training breaks.

loss.item() converts a one-element tensor to a Python float. Accumulating the tensor itself would keep the entire computation graph in memory.

Step 6: Evaluate the Model

def evaluate(model, loader, loss_fn):
    model.eval()
    correct, total_loss = 0, 0.0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            logits = model(images)
            total_loss += loss_fn(logits, labels).item() * images.size(0)
            correct += (logits.argmax(dim=1) == labels).sum().item()
    n = len(loader.dataset)
    return total_loss / n, correct / n

Two details matter:

• model.eval() switches layers like Dropout and BatchNorm to inference behavior. Our MLP has neither, but build the habit now.
• torch.no_grad() disables gradient tracking, cutting memory use and speeding up inference.

Step 7: Run Training

EPOCHS = 5
for epoch in range(1, EPOCHS + 1):
    train_loss = train_one_epoch(model, train_loader, loss_fn, optimizer)
    test_loss, test_acc = evaluate(model, test_loader, loss_fn)
    print(f"Epoch {epoch}: train_loss={train_loss:.4f} "
          f"test_loss={test_loss:.4f} test_acc={test_acc:.2%}")

Expect output similar to this (exact numbers vary by seed and hardware):

Set a seed with torch.manual_seed(42) before building the model if you need reproducible runs.

Step 8: Save, Load, and Predict

Save the state dict, not the whole model object. It is portable and robust to code refactors.

torch.save(model.state_dict(), "mnist_mlp.pt")

# Later:
model = MLP().to(device)
model.load_state_dict(torch.load("mnist_mlp.pt", map_location=device))
model.eval()

image, label = test_data[0]
with torch.no_grad():
    pred = model(image.unsqueeze(0).to(device)).argmax(1).item()
print(f"Predicted: {pred}, Actual: {label}")

unsqueeze(0) adds the batch dimension, since the model expects shape (B, 1, 28, 28).

Common Beginner Bugs and Fixes

A loss near 2.30 equals ln(10), the loss of random guessing across 10 classes. If your loss never moves below it, the model is not learning.

Real-World Workflows: Where to Go Next

Upgrade to a Convolutional Network

MLPs ignore spatial structure. A small CNN usually pushes MNIST accuracy to 99%+:

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 14x14
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(), nn.Dropout(0.25),
            nn.Linear(64 * 7 * 7, 128), nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

Swap MLP() for CNN(). Nothing else in the training loop changes. That is the payoff of PyTorch’s modular design.

Speed It Up With torch.compile

PyTorch 2.x can compile your model into optimized kernels:

model = torch.compile(model)

Gains depend on model size and hardware. Small MLPs see little benefit, while larger models often train noticeably faster.

Production Use Cases

Framework Comparison: Why PyTorch for Beginners?

PyTorch’s eager execution means you can drop a print() or breakpoint anywhere in forward and inspect real values. That makes it the easiest framework to learn by experimentation.

Frequently Asked Questions

How long does it take to train a first neural network in PyTorch?

On a modern laptop CPU, the MNIST MLP above trains in roughly 1–3 minutes for 5 epochs. A GPU cuts that to seconds. The dataset is small, so hardware is rarely a bottleneck.

Do I need a GPU to learn PyTorch?

No. Small models like MLPs and basic CNNs run fine on a CPU. For larger projects, free tiers on Google Colab and Kaggle provide GPU access.

What is the difference between nn.Module and nn.Sequential?

nn.Sequential chains layers in a fixed order. nn.Module is the base class for any model and lets you write custom forward logic such as skip connections, branching, or multiple inputs. Use Sequential for simple stacks and Module for everything else.

Why is my PyTorch model not learning?

Check four things in order: gradients are zeroed each step, the learning rate is reasonable (1e-3 for Adam), labels are integer class indices for CrossEntropyLoss, and the data is normalized. Also confirm that loss is not stuck near ln(num_classes).