Training a Simple Neural Network in PyTorch and Integrating with Gradio for MNIST Digit Recognition
Tags: ai, mnist, pytorch
Published: 2023-05-23 17:35:39

---
title: Training a Simple Neural Network in PyTorch and Integrating with Gradio for MNIST Digit Recognition
published: true
description: 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
tags: #ai #mnist #pytorch
# cover_image: https://direct_url_to_image.jpg
# Use a ratio of 100:42 for best results.
# published_at: 2023-05-11 11:51 +0000
---


## 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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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

```python
!pip install gradio
```

### Step 2: Import Gradio and define the prediction function

```python
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:

```python
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:

```python
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.