Building Your First Neural Network with PyTorch

Let’s build our first neural network with pytorch and try to understand how each part works. We will build a neural network to classify handwritten digits from MNIST dataset.

Setting up the environment

First, we have to install pytorch and torchvision.

pip install torch torchvision

I am using google colab, where pytrch is pre-installed and we also have access to GPUs.

Understanding the dataset

MNIST dataset, a classic benchmark of 60,000 training and 10,000 test images of handwritten digits (0–9), each 28×28 pixels.

Here’s how to load it:

import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([transforms.ToTensor()])
trainset = torchvision.datasets
    .MNIST(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets
    .MNIST(root='./data', train=False, download=True, transform=transform)
trainloader = torch.utils.data
    .DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data
    .DataLoader(testset, batch_size=64, shuffle=False)

transforms.toTensor() converts the python imaging library(PIL) to pytorch tensors, scaling the pixel values to the range [0, 1]. This will be used to transforms the images from MNIST to tensors.