Training a Simple Neural Network in PyTorch and Integrating with Gradio for MNIST Digit Recognition

Introduction

In this article, we will walk through the steps of training a simple neural network on the MNIST dataset using PyTorch and then deploying it with Gradio for interactive predictions. The MNIST dataset is a popular dataset in the field of machine learning that consists of 70,000 28x28 grayscale images of handwritten digits.

Training a Neural Network with PyTorch

PyTorch is an open-source deep learning framework developed by Facebook's artificial intelligence research group. It provides a wide range of functionalities for building and training neural networks.

Step 1: Import necessary libraries

First, we need to import PyTorch, torchvision (a package with popular datasets, model architectures, and common image transformations), and some other necessary libraries:

import torch import torchvision import torchvision.transforms as transforms from torch import nn, optim 

Step 2: Load the dataset

Next, we load the MNIST dataset. We'll use torchvision's built-in functionality to do this. We also apply transformations to normalize the data:

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

Step 3: Define the network

We'll define a simple feed-forward neural network with one hidden layer:

class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28 * 28, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = x.view(-1, 28 * 28) x = torch.relu(self.fc1(x)) x = self.fc2(x) return x net = Net() 

Step 4: Define the loss function and optimizer

We'll use CrossEntropyLoss for our loss function and SGD for our optimizer:

criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9) 

Step 5: Train the network

Now we're ready to train our network:

for epoch in range(10): running_loss = 0.0 for i, data in enumerate(trainloader, 0): inputs, labels = data optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch {epoch+1}, loss: {running_loss/len(trainloader)}') print('Finished Training') 

Deploying with Gradio

Gradio is an open-source library for creating customizable UI components around your ML models. It allows us to demonstrate a model’s functionality in an intuitive manner.

Step 1: Install Gradio

!pip install gradio 

Step 2: Import Gradio and define the prediction function

import gradio as gr def predict(image): image = image.reshape(1, 1, 28, 28) image = torch.from_numpy(image).float() output = net(image) _, predicted = torch.max(output.data, 1) return predicted.item() 

In the predict function, we take the input image, reshape it to match our model's expected input shape, convert it to a torch tensor, pass it through our model to get the output, and then return the predicted digit.

Step 3: Define the Gradio interface

Now, we define the interface for our model. We'll use an 'Image' input interface and a 'Label' output interface:

iface = gr.Interface( fn=predict, inputs=gr.inputs.Image(shape=(28, 28), invert_colors=True, source="canvas"), outputs="label", interpretation="default" ) 

The 'Image' input interface lets users draw an image with their mouse. We set invert_colors=True because the MNIST dataset consists of white digits on a black background, and by default, the Gradio image interface has a white background.

Step 4: Launch the interface

Finally, we launch the interface:

iface.launch() 

With this, you should see an interactive interface where you can draw a digit and see the prediction from your PyTorch model.

Conclusion

In this article, we saw how to train a simple neural network using PyTorch and then deploy it with Gradio for interactive predictions. This combination allows us to leverage the power of deep learning models in an easy-to-use and interpret manner.