PyTorch MNIST: Neural Network (Part 1) – Training in PyTorch and NLP #2

Introduction:

Welcome to part 1 of the PyTorch MNIST Neural Network tutorial series. In this tutorial, we will walk you through the process of building a simple neural network using PyTorch to classify hand-written digits from the MNIST dataset. We will cover the basics of PyTorch, as well as how to create a neural network model and train it using the popular MNIST dataset. By the end of this tutorial, you should have a good understanding of how to build and train your own neural network using PyTorch.

What is PyTorch?

PyTorch is an open-source deep learning library developed by Facebook AI Research. It is widely used for building deep neural networks and conducting research in the field of artificial intelligence. PyTorch provides a flexible and intuitive platform for developing deep learning models, with a focus on simplicity and ease of use. It also supports popular deep learning techniques such as automatic differentiation, GPU acceleration, and distributed computing.

MNIST Dataset:

The MNIST dataset is a widely used dataset in the field of deep learning for training and testing machine learning models. It consists of a set of 60,000 training images and 10,000 test images of hand-written digits from 0 to 9. Each image is a 28×28 pixel grayscale image, making it a popular choice for training image classification models.

Building a Neural Network Model:

Now, let’s start by building a simple neural network model using PyTorch. We will create a fully connected neural network with one hidden layer and the following architecture:

• Input layer: 28×28 = 784 neurons (flattened image)
• Hidden layer: 128 neurons
• Output layer: 10 neurons (corresponding to 10 classes – digits 0 to 9)

First, we need to import the necessary libraries:

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

Next, let’s define the neural network model:

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

In the __init__ method, we define the layers of our neural network model. We use a nn.Sequential container to define a sequence of layers consisting of two fully connected layers with ReLU activation functions. The forward method specifies how the input data is passed through the layers of the model to generate the output logits.

Preparing the MNIST Dataset:

Next, let’s prepare the MNIST dataset for training our neural network model. We will use the torchvision library to load and preprocess the dataset.

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

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

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

In the code above, we define a set of transformations to be applied to the input images, including converting them to tensors and normalizing the pixel values. We then load the training and test datasets using the torchvision.datasets.MNIST class and create data loaders to iterate over the dataset in batches during training.

Training the Neural Network Model:

Now, let’s train our neural network model on the MNIST dataset using PyTorch. We will define the loss function, optimizer, and training loop to optimize the model parameters.

model = NeuralNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

def train_model(model, criterion, optimizer, train_loader, num_epochs):
    for epoch in range(num_epochs):
        model.train()
        for inputs, targets in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')

train_model(model, criterion, optimizer, train_loader, num_epochs=5)

In the code above, we define the training loop train_model that iterates over the training dataset for a specified number of epochs. During each epoch, we pass the input images through the neural network model to compute the output logits, calculate the cross-entropy loss, backpropagate the gradients, and update the model parameters using stochastic gradient descent.

Conclusion:

In this tutorial, we have explored the basics of building and training a simple neural network model using PyTorch on the MNIST dataset. We have covered how to define the neural network architecture, prepare the dataset, and train the model using PyTorch functionalities. In part 2 of this tutorial series, we will evaluate the trained model on the test set and visualize the results. Stay tuned for the next part!