Build and deploy Image Classification web app using Django on Heroku (Part-1)
“No machine learning model is valuable, unless it’s deployed to production.” — Luigi Patruno
In this article, we will see how to
- Build a SOTA model on Image Classification on CIFAR100 dataset using SAM Optimizer (Part-1)
- Create a web app using Django to predict the class of the image (Part-2)
- Deploy the web app on Heroku using GitHub (Part-3)
In Part-1, lets focus on building the model using PyTorch. An universal workflow for addressing any machine learning problem is as follows:
Lets get started by importing the required libraries
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torchvision
import torchvision.transforms as transforms
import math
import os
import argparseChecking for the availability of GPU else using CPU as ‘device’ and initializing some settings
device = 'cuda' if torch.cuda.is_available() else 'cpu'
lr = 0.1 # learning rate
best_acc = 0
num_epochs = 225Setting the transformations to be applied on the Train and Test dataset
train_transforms = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])test_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914,0.4822,0.4465),(0.2023,0.1994,0.2010)),
])
- Collecting the Data: Gathering the CIFAR 100 data from the torchvision datasets
- Preparing the Data: Tranforms are applied on the train and test data
train_data = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=train_transforms)train_loader = torch.utils.data.DataLoader(train_data, batch_size=128, shuffle=True, num_workers=0)
test_data = torchvision.datasets.CIFAR100(root='./data', train = False, download=True, transform=test_transforms)test_loader = torch.utils.data.DataLoader(test_data, batch_size=128, shuffle=False, num_workers=0)
3. Creating Model
net = torchvision.models.resnet18(pretrained=True)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs,100)
net = net.to(device)The commonly done optimization of only training loss will lead to a suboptimal model quality. SAM [Sharpness Aware Minimization] is a novel effective procedure for minimizing loss value and loss sharpness simultaneously. It seeks parameters that lie in neighborhood having uniformly low loss.
class SAM(torch.optim.Optimizer): def __init__(self, params, base_optimizer, rho=0.05, **kwargs): assert rho >= 0.0, f”Invalid rho, should be non-negative: {rho}”
defaults = dict(rho=rho, **kwargs)
super(SAM, self).__init__(params, defaults)
self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
self.param_groups = self.base_optimizer.param_groups @torch.no_grad()
def first_step(self, zero_grad=False):
grad_norm = self._grad_norm()
for group in self.param_groups:
scale = group[“rho”] / (grad_norm + 1e-12)
for p in group[“params”]:
if p.grad is None: continue
e_w = p.grad * scale
p.add_(e_w) # climb to the local maximum “w + e(w)”
self.state[p][“e_w”] = e_w if zero_grad:
self.zero_grad() @torch.no_grad()
def second_step(self, zero_grad=False):
for group in self.param_groups:
for p in group[“params”]:
if p.grad is None: continue
p.sub_(self.state[p][“e_w”]) # get back to “w” from “w + e(w)” self.base_optimizer.step() # do the actual “sharpness-aware” update
if zero_grad: self.zero_grad() def step(self, closure=None):
raise NotImplementedError(“SAM doesn’t work like the other optimizers, you should first call `first_step` and the `second_step`; see the documentation for more info.”) def _grad_norm(self):
norm = torch.norm(torch.stack([
p.grad.norm(p=2)
for group in self.param_groups
for p in group[“params”]
if p.grad is not None ]),
p=2) return norm
Initializing Loss function, optimizer and scheduler
criterion = nn.CrossEntropyLoss()base_optimizer = torch.optim.SGDoptimizer = SAM(net.parameters(), base_optimizer, rho=0.05, lr=0.1, momentum=0.9, weight_decay=5e-4)scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
4. Training the model
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
H = []
Y = []
for batch_idx, (inputs, targets) in enumerate(train_loader):
Y.append(targets)
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.mean().backward()
optimizer.first_step(zero_grad=True) criterion(net(inputs), targets).mean().backward()
optimizer.second_step(zero_grad=True)
train_loss += loss.mean().item() _, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item() if(batch_idx % 50 == 0):
print(str(batch_idx)+"/"+str(len(train_loader)) +" Loss: " + str(train_loss/(batch_idx+1)) +" Acc: "+ str(100.*correct/total)) H.append(outputs.to('cpu')) H = torch.cat(H,0).detach().numpy()
Y= torch.cat(Y,0).numpy() return train_loss / len(train_loader), 100.*correct/total, H, Y
5. Testing the model
def test(epoch):
global best_acc
net.eval()
test_loss = 0
correct = 0
total = 0
H = []
Y = [] with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(test_loader):
Y.append(targets)
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.mean().item()
_, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() if(batch_idx % 50 == 0):
print(str(batch_idx)+"/"+str(len(test_loader)) +" Loss: " + str(test_loss/(batch_idx+1)) +" Acc: "+ str(100.*correct/total)) H.append(outputs.to('cpu')) acc = 100.*correct/total if acc > best_acc:
print('Saving..')
state = { 'net': net.state_dict(),'acc': acc, 'epoch': epoch,} if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint') torch.save(state, './checkpoint/ckpt.pth')
best_acc = acc H = torch.cat(H,0).detach().numpy()
Y = torch.cat(Y,0).numpy() return test_loss / len(test_loader), 100.*correct/total, H, Y
Finally, starting the training and testing loops and saving the best model
train_loss = []
train_accuracy = []
test_loss = []
test_accuracy = []for epoch in range(0, num_epochs):
tr_loss, tr_acc, _, _ = train(epoch)
te_loss, te_acc, _, _ = test(epoch) train_loss.append(tr_loss)
train_accuracy.append(tr_acc)
test_loss.append(te_loss)
test_accuracy.append(te_acc) print("Train Accuracy: " , max(train_accuracy), "% ", " Test Accuracy: " , max(test_accuracy), "%") scheduler.step()
The best model will be stored in the checkpoint folder
Find the code at:
