Image Classification Using Pytorch Lightning
Introduction
In this tutorial, I’ll be using PyTorch lightning to classify Images from the https://www.kaggle.com/brsdincer/vehicle-detection-image-set Dataset, the aim is to classify images into vehicles and non-vehicles using PyTorch lightning. This tutorial assumes you are familiar with Data Science and Pytorch.
The Dataset
The Dataset contains two folders containing images of vehicles and non-vehicles, respectively. We are tasked with creating a classifier that can classify vehicles and non-vehicles.
The Notebook
This Notebook was run on a Kaggle notebook using a GPU. Note that running this notebook without a GPU might take a longer time. The original Notebook can be found here
Dependencies
If running the cell below gives you errors, then one or more Modules haven’t been installed in your environment.
import pandas as pd
import os
from sklearn.model_selection import train_test_split
import globimport torch
from torch import nn
import torchvision
from torch.utils.data import Dataset,DataLoader
import cv2
import pytorch_lightning as pl
import torchvision.transforms as transforms
Preprocessing
Next, we’ll be making a Dataframe containing paths and targets to make accessing the data easier.
vehicles = glob.glob(f"../input/vehicle-detection-image-set/data/vehicles/*.png")#returns a list of paths of images in Vehicles foldernon_vehicles = glob.glob(f"../input/vehicle-detection-image-set/data/non-vehicles/*.png")#returns a list of paths of images in non-Vehicles folderdata1 = {
"mask_id": vehicles,
"mask":"vehicles"
}
data2 = {
"mask_id": non_vehicles,
"mask":"non_vehicles"
}
vehicles = pd.DataFrame(data1)# a dataframe with the path and labels
non_vehicles = pd.DataFrame(data2)df = pd.concat([vehicles,non_vehicles],ignore_index= True)#combining the two datasets
dfdef label(x):
if x == "vehicles":
return 1
if x == "non_vehicles":
return 0df["mask"] = df["mask"].apply(label)# encoding the label numerically
df
Next, We’ll make a Pytorch-Lightning DataModule to send to the trainer
X = df["mask_id"]
y = df["mask"]
x_train , x_test , y_train, y_test = train_test_split(X,y,test_size = 0.25,random_state = 42 ,)
x_val , x_test , y_val, y_test = train_test_split(x_test,y_test,test_size = 0.25,random_state = 42 ,)train_dataset = MaskDataset(
x_train.values,
y_train.values
)test_dataset = MaskDataset(
x_test.values,
y_test.values
)
prediction_dataset = MaskDataset(
x_test.values,
)validation_dataset = MaskDataset(
x_val.values,
y_val.values
)class Mask_DataModule(pl.LightningDataModule):
def __init__(self):
super().__init__()
self.train = train_dataset
self.val = validation_dataset
self.test = test_dataset
self.prediction = prediction_dataset
def train_dataloader(self):
return DataLoader(self.train,batch_size = 64,shuffle = True,num_workers=2)
def val_dataloader(self):
return DataLoader(self.val,batch_size = 32,shuffle = False,num_workers=1)
def test_dataloader(self):
return DataLoader(self.test,batch_size = 32,shuffle = False,num_workers=1)
def predict_dataloader(self):
return DataLoader(self.prediction,batch_size = 1,shuffle = False,num_workers=1)
The Model
We’ll be using a Resnet50 model. If you're using your own neural network, then it might take a longer training period to get the same results.
neural_network = torchvision.models.resnet50(pretrained= True)
neural_network.fc = torch.nn.Linear(2048,2)# changing the number of output features to 2Next, we’ll use the Lightning Module to make the model. We’ll be using StepLr scheduler for learning rate
loss_func = torch.nn.CrossEntropyLoss()class Vehicle_Model(pl.LightningModule):
def __init__(self):
super().__init__()
self.neural_net = neural_network
def forward(self,x):
return self.neural_net(x)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters() , lr = 1e-4)
#choosing a optimizer
sch = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 20, gamma=0.5, last_epoch=-1, verbose=False)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": sch,
"monitor": "val_loss",
},
}
def training_step(self,batch,batch_idx):
x,y = batch
y_pred = self(x)
#y = y.unsqueeze(-1)
loss = loss_func(y_pred,y)
return loss
def validation_step(self,batch,batch_idx):
x,y = batch
y_pred = self(x)
#y = y.unsqueeze(-1)
loss = loss_func(y_pred,y)
self.log("val_loss" ,loss )
return loss
def test_step(self,batch,batch_idx):
x,y = batch
y_pred = self(x)
#y = y.unsqueeze(-1)
loss = loss_func(y_pred,y)
self.log("test_loss : " , loss)
return lossfrom pytorch_lightning.callbacks import ModelCheckpoint
checkpoint_callback = ModelCheckpoint(
monitor = "val_loss",#monitors val loss
mode = "min",#Picks the fold with the lowest val_loss
)
The model can be trained using one epoch if you’ve followed the tutorial so far. The best part about PyTorch lightning is that you can set the number of gpus by simply setting “ gpus = [number of gpus]“
%%time # Checking the amount of time the cell takes to run
from pytorch_lightning import Trainer
model = Vehicle_Model()
module = Vehicle_DataModule()
trainer = Trainer(max_epochs=1,gpus = 1,callbacks = [checkpoint_callback])
trainer.fit(model,module)
Testing the Model
trainer.test()
Generating the Predictions
predictons = trainer.predict()
probs = nn.Softmax() # since the outputs are logits we use the softmax function to obtain probabilitiesdef get_score(y_pred,y):
probabilities = []
accuracy = 0
for x in y_pred:
prob = probs(x)
top_p, top_class = prob.topk(1, dim = -1)
probabilities.append(float(top_class))
i = 0
for x in y:
if int(x) == int(probabilities[i]):
accuracy = accuracy + 1
i = i +1
score = accuracy/len(probabilities)
return score # returns accuracyget_score(predictons,y_test.values)
As can be seen we obtain a accuracy of 97.9%
