# Building Your First Neural Network with PyTorch

> Build and train your first PyTorch neural network end‑to‑end: tensors, nn.Module, training loops, metrics — perfect for beginners.

Published: 2026-03-01  
Canonical: https://aimastermindlab.com/tutorials/first-neural-network-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.

```python
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):

```python
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.

```python
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.

```python
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).

```python
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.
---

Source: https://aimastermindlab.com/tutorials/first-neural-network-pytorch — AI Mastermind Lab (https://aimastermindlab.com)