# Image Classification with PyTorch Lightning (Simple ANN)

# Introduction

Hello everyone, welcome to my Blog in this tutorial, we will be learning how to do simple image classification as we did in the [last tutorial](https://www.chalisebibek.com.np/pytorch-101-image-classification-with-simpleann) but this time, we will be leveraging PyTorch Lightning to do the same task but with ease. So, I am planning to release both core PyTorch implementation as well as PyTorch Lightning implementation.

![Alt text](https://www.datanami.com/wp-content/uploads/2023/03/lightning-ai.png align="left")

Without much due, let's dive directly into the code.

`In Local Machine`

```bash
    pip3 install lightning
```

`In Colab Notebook`

```python
    !pip install lightning
```

At first, we need to install `lightning` in our device. For this tutorial, we will be using `Colab`, and in colab also we have to install it.

```python
import os
import torch
from torch import nn
import torchvision
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader
import lightning.pytorch as pl
```

So, all the required modules are imported. Everything looks almost the same, just the import of `lightning.pytorch` as pl is new. Then, let's copy a few more lines from previous tutorials, nothing new, the same lines and the same function.

```python

# Data preprocessing and augmentation
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# Load CIFAR-10 dataset
dataset = torchvision.datasets.CIFAR10(root='./cifar', train=True, transform=transform, download=True)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)

criterion = nn.CrossEntropyLoss()
```

The model definition will also be the same.

```python
class SimpleANN(nn.Module):
    def __init__(self):
        super(SimpleANN, self).__init__()
        self.fc1 = nn.Linear(3 * 32 * 32, 128)  # Flatten the 32x32 RGB images
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)  # Output layer for 10 classes
    
    def forward(self, x):
        x = x.view(-1, 3 * 32 * 32)  # Flatten the input images
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
```

```python
class LitSimpleANN(pl.LightningModule):
  def __init__(self, model):
    super().__init__()
    self.model = model
  
  def training_step(self, batch, batch_idx):
    x, y = batch
    # print(y)
    out = self.model(x)
    # print(out)
    loss = criterion(out, y)
    preds = torch.argmax(out, dim=1)
    accuracy = torch.sum(preds == y).item() / y.size(0)
    
    self.log('train_loss', loss,  on_epoch=True)  # Log training loss
    self.log('train_accuracy', accuracy,  on_epoch=True)
    return loss
  
  def configure_optimizers(self):
    optimizer = torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
    return optimizer
```

Now, the real magical lines of PyTorch lightning come into action. We defined a class `LitSimpleANN` as a subclass of `PyTorch Lightning Module` and this class takes the above-defined `SimpleANN` class as input. Then, we define the training step inside this `LitSimpleANN` class. `training_step` function takes `batch` and `batch_idx` as input. So each batch has input images and the labels associated with the batch. So, Same as before, we fed the model the input image, it gives the probability distribution of each class and we calculate loss and accuracy on it. Here, `self.log` stores `loss` and `accuracy` per `step` and per `epoch` which can be useful during the evaluation of the model performance.

```python
model = LitSimpleANN(SimpleANN())
```

### Training the data

```python
from pytorch_lightning.loggers.csv_logs import CSVLogger

trainer = pl.Trainer(max_epochs=10,  logger=CSVLogger("logs"))
trainer.fit(model = model, train_dataloaders=dataloader)
```

Here, I have imported CSVLogger to save the logs in `CSV` format for further reference. Here, `pl.Trainer` takes away all the manual training loops from the custom PyTorch training loop we also discussed in the last video. So, when we hit shift enter, `Bam`, the training starts. Nothing much fancy here, `max_epochs=10` is for the number of epochs, as the default is `1000 epochs`. The output of this code is as follows:

![](https://cdn.hashnode.com/res/hashnode/image/upload/v1693290932840/ed3aef5f-799e-44f3-a581-231669728fa8.png align="left")

Let's plot some graphs.

```python
import pandas as pd
import seaborn as sn
metrics = pd.read_csv(f"{trainer.logger.log_dir}/metrics.csv")
del metrics["step"]
metrics.set_index("epoch", inplace=True)
sn.relplot(data=metrics, kind="line")
```

The `mertics.csv` is auto-generated by trainer logger and this will give a nice line graph.

![](https://cdn.hashnode.com/res/hashnode/image/upload/v1693290919549/1178dcc1-f3d9-4212-a6c1-6f563f1a4c80.png align="center")

## Conclusion

In this post, we learned how to implement image classification using simple ANN in the PyTorch Lightning library. Lightening AI is fast, easy to use and saves us from unnecessary loops of code. I highly recommend you all learn about it and it will save you from hassle compared to core PyTorch code.

This much from today, so see you guys in the next [tutorial](https://www.chalisebibek.com.np/a-slightly-advanced-ann-in-pytorch-for-image-classification?source=more_series_bottom_blogs).

[Colab](https://colab.research.google.com/drive/1ajNmMB1OyvMYlcWdKOuV0tk1YewhINtP?usp=sharing)

[YouTube](https://www.youtube.com/@bibekchalise)

`Bibek Chalise is a Machine Learning enthusiast, Computer Vision Scientist and is associated with MarginTop Solutions.`

`MarginTop Solutions`

`Where Tech Meets Brilliance`

`Pokhara, Nepal`

`margintopsolutions@gmail.com`

[`Website`](https://margintopsolutions.github.io)

![](https://cdn.hashnode.com/res/hashnode/image/upload/v1693290740053/c2cefc57-623d-466a-8eff-c10de83301af.png align="left")