Skip to content

Feed Forward Neural Network - PyTorch Beginner 13

In this part we will implement our first multilayer neural network that can do digit classification based on the famous MNIST dataset.

Learn all the basics you need to get started with this deep learning framework! In this part we will implement our first multilayer neural network that can do digit classification based on the famous MNIST dataset.

We put all the things from the last tutorials together:

  • Use the DataLoader to load our dataset and apply a transform to the dataset
  • Implement a feed-forward neural net with input layer, hidden layer, and output layer
  • Apply activation functions.
  • Set up loss and optimizer
  • Training loop that can use batch training.
  • Evaluate our model and calculate the accuracy.
  • Additionally, we will make sure that our whole code can also run on the gpu if we have gpu support.

All code from this course can be found on GitHub.

Neural Network in PyTorch

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Hyper-parameters 
input_size = 784 # 28x28
hidden_size = 500 
num_classes = 10
num_epochs = 2
batch_size = 100
learning_rate = 0.001

# MNIST dataset 
train_dataset = torchvision.datasets.MNIST(root='./data', 

test_dataset = torchvision.datasets.MNIST(root='./data', 

# Data loader
train_loader =, 

test_loader =, 

examples = iter(test_loader)
example_data, example_targets =

for i in range(6):
    plt.imshow(example_data[i][0], cmap='gray')

# Fully connected neural network with one hidden layer
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.input_size = input_size
        self.l1 = nn.Linear(input_size, hidden_size) 
        self.relu = nn.ReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)  

    def forward(self, x):
        out = self.l1(x)
        out = self.relu(out)
        out = self.l2(out)
        # no activation and no softmax at the end
        return out

model = NeuralNet(input_size, hidden_size, num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  

# Train the model
n_total_steps = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):  
        # origin shape: [100, 1, 28, 28]
        # resized: [100, 784]
        images = images.reshape(-1, 28*28).to(device)
        labels =

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize

        if (i+1) % 100 == 0:
            print (f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')

# Test the model
# In test phase, we don't need to compute gradients (for memory efficiency)
with torch.no_grad():
    n_correct = 0
    n_samples = 0
    for images, labels in test_loader:
        images = images.reshape(-1, 28*28).to(device)
        labels =
        outputs = model(images)
        # max returns (value ,index)
        _, predicted = torch.max(, 1)
        n_samples += labels.size(0)
        n_correct += (predicted == labels).sum().item()

    acc = 100.0 * n_correct / n_samples
    print(f'Accuracy of the network on the 10000 test images: {acc} %')

FREE VS Code / PyCharm Extensions I Use

鉁 Write cleaner code with Sourcery, instant refactoring suggestions: Link*

PySaaS: The Pure Python SaaS Starter Kit

馃殌 Build a software business faster with pure Python: Link*

* These are affiliate link. By clicking on it you will not have any additional costs. Instead, you will support my project. Thank you! 馃檹