1. Import libraries

import torch
import torch.nn as nn
import torch.optim as optim

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

torch → main PyTorch library
nn → neural-network layers
optim → optimizers like Adam
DataLoader → creates batches
datasets → provides MNIST
transforms → image preprocessing

2. Select device

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

Uses GPU if available.
Otherwise uses CPU.

model.to(device)

moves the model to that device.

3. Transform images

transform = transforms.ToTensor()

Converts images into PyTorch tensors. MNIST pixel values are converted roughly from:

0–255 → 0–1

4. Load MNIST

train_dataset = datasets.MNIST(
    root="./data",
    train=True,
    download=True,
    transform=transform
)

train=True → training data
MNIST training set → 60,000 images

test_dataset = datasets.MNIST(
    root="./data",
    train=False,
    download=True,
    transform=transform
)

train=False → testing data
MNIST test set → 10,000 images

5. Create DataLoader

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True
)

batch_size=64 → process 64 images at a time
shuffle=True → randomly mix training data

test_loader = DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False
)

Testing doesn’t need shuffling.

6. MNIST image

Each MNIST image is:

28 × 28

So total pixels:

28 × 28 = 784

Image tensor:

1 × 28 × 28

1 means grayscale channel.

7. Neural network

class SimpleNN(nn.Module):

Creates a custom neural network.

First layer

self.fc1 = nn.Linear(28*28, 128)

Means:

784 inputs → 128 neurons

Second layer

self.fc2 = nn.Linear(128, 10)

Means:

128 neurons → 10 outputs

Why 10? Because MNIST has:

0 1 2 3 4 5 6 7 8 9

8. Forward pass

def forward(self, x):

    x = x.view(x.size(0), -1)

    x = self.fc1(x)

    x = torch.relu(x)

    x = self.fc2(x)

    return x

Flow:

28 × 28 image
      ↓
Flatten
      ↓
784
      ↓
Linear layer
      ↓
128
      ↓
ReLU
      ↓
10 outputs

9. Flatten

x.view(x.size(0), -1)

Converts:

1 × 28 × 28

into:

For a batch of 64:

64 × 1 × 28 × 28
        ↓
64 × 784

10. ReLU

torch.relu(x)

Formula:

ReLU(x) = max(0, x)

Examples:

-5 → 0
-2 → 0
 3 → 3
 7 → 7

ReLU adds non-linearity, allowing the network to learn complex patterns.

11. Create model

model = SimpleNN().to(device)

Creates the neural network and moves it to CPU/GPU.

12. Loss function

criterion = nn.CrossEntropyLoss()

Loss tells us:

How wrong is the prediction?

For MNIST classification, CrossEntropyLoss is commonly used.

Low loss  → good prediction
High loss → bad prediction

13. Optimizer

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

Adam updates the model’s weights.

Learning rate

lr=0.001

Controls how big each weight update is.

14. Epoch

epochs = 10

One epoch means:

The model has seen the entire training dataset once.

So:

10 epochs = 10 complete passes

15. Training loop

The most important part:

optimizer.zero_grad()

outputs = model(images)

loss = criterion(outputs, labels)

loss.backward()

optimizer.step()

Remember:

1. Clear gradients
       ↓
2. Make prediction
       ↓
3. Calculate loss
       ↓
4. Backpropagation
       ↓
5. Update weights

16. `optimizer.zero_grad()`

optimizer.zero_grad()

Clears old gradients before calculating new ones.

17. Forward pass

outputs = model(images)

The images go through the neural network.

Image
 ↓
784
 ↓
128
 ↓
ReLU
 ↓
10 scores

18. Calculate loss

loss = criterion(outputs, labels)

Compares:

Prediction
   vs
Correct label

19. Backpropagation

loss.backward()

Calculates gradients. Gradients tell the model how its weights should change to reduce the loss.

20. Update weights

optimizer.step()

Adam uses the gradients to update the weights. This is where the model learns.

21. Training vs Testing

Training

model.train()

Used when learning. The model:

Input
 ↓
Prediction
 ↓
Loss
 ↓
Backward
 ↓
Update weights

Testing

model.eval()

Used after training to measure performance. No weight updates happen.

22. `torch.no_grad()`

with torch.no_grad():

During testing, we don’t need gradients. Benefits:

Uses less memory
Faster computation
No training happens

23. Prediction

_, predicted = torch.max(outputs, 1)

The model produces 10 scores. Example:

→ 0.2
→ 1.1
→ 0.4
→ 7.8  ← highest
→ 0.3
...

Highest score is class 3. Therefore:

Prediction = 3

24. Accuracy

accuracy = 100 * correct / total

Example:

Correct = 9,800
Total   = 10,000

Then:

Accuracy = 98%

25. Important correction

Your original code has:

print("Test Accuracy: {accuracy: .2f}%")

This is incorrect. Use:

print(f"Test Accuracy: {accuracy:.2f}%")

The f makes it an f-string, so Python replaces {accuracy} with its value.

Revision

MNIST
 ↓
60,000 training images
10,000 testing images
 ↓
ToTensor()
 ↓
28 × 28 image
 ↓
Flatten
 ↓
784 pixels
 ↓
Linear(784 → 128)
 ↓
ReLU
 ↓
Linear(128 → 10)
 ↓
10 class scores
 ↓
CrossEntropyLoss
 ↓
loss.backward()
 ↓
Adam optimizer
 ↓
Update weights
 ↓
Repeat for 10 epochs
 ↓
Test model
 ↓
Calculate accuracy

Most important 5 lines:

outputs = model(images)       # Prediction

loss = criterion(outputs, labels)  # Error

loss.backward()               # Calculate gradients

optimizer.step()              # Update weights

optimizer.zero_grad()         # Clear old gradients

In one sentence: The neural network takes a 28×28 handwritten digit, converts it into 784 numbers, processes them through 128 neurons, produces 10 class scores, calculates the error, and repeatedly updates its weights until its predictions improve.

​1. Import libraries

​2. Select device

​3. Transform images

​4. Load MNIST

​5. Create DataLoader

​6. MNIST image

​7. Neural network

​First layer

​Second layer

​8. Forward pass

​9. Flatten

​10. ReLU

​11. Create model

​12. Loss function

​13. Optimizer

​Learning rate

​14. Epoch

​15. Training loop

​16. optimizer.zero_grad()

​17. Forward pass

​18. Calculate loss

​19. Backpropagation

​20. Update weights

​21. Training vs Testing

​Training

​Testing

​22. torch.no_grad()

​23. Prediction

​24. Accuracy

​25. Important correction

​Revision

​Most important 5 lines: