Transfer Learning with Pytorch

Transfer learning is a powerful way to solve overfitting issue related to small dataset. Pytorch Training is a powerful deep learning framework to implement tr ansfer learning . In this article, we will demonstrate how to apply transfer technique to solve the overfitting issue with Pytorch

First we load the small dataset of bees and ants images to a. trainloader.

mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]

transforms = {
    'train': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ]),
}

data_dir = '/content/drive/MyDrive/dataset/hymenoptera_data'

trainset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transforms['train'])
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)

testset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transforms['val'])
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=True)

Then we load the pre-trained model from torchvision models library, for example vgg16 model. You can use any pre-trained model available on Pytrch.

import torchvision.models as models

model = models.vgg16(pretrained=True)

Then we freeze all the features parameters in the model as we plan to reuse the trained features of the pretrained model. Next we replace the classifier since we only want to classify ants and bees instead of 1000 imagenet classes. We also assign the model to GPU.

for param in model.parameters():
  param.requires_grad = False

model.classifier[-1] = nn.Sequential(
                       nn.Linear(in_features=4096, out_features=10)
                        )

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
model.to(device)

Next we define the loss function and optimizer as follows:

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
model.to(device)

Next we re-train the model as follows:

from tqdm.notebook import trange,tqdm

num_epochs = 2

for i in trange(num_epochs):
    train_loss = 0
    
    for (X, y) in tqdm(trainloader):
        X = X.to(device)
        y = y.to(device)

        yhat = model(X)
        
        optimizer.zero_grad()
        loss = criterion(yhat, y)
        loss.backward()
        optimizer.step()
        
        train_loss  += loss.item()
    
    print(f"Epoch({i+1}/{num_epochs}) | Training loss: {train_loss/len(trainloader)} | ",end='')

We can evaluate the model as follows

from tqdm.notebook import trange,tqdm

model.eval()
num_correct = 0
total = 0

for (X,y) in tqdm(testloader):
    X = X.to(device)
    y = y.to(device)
    
    yhat = model(X)
    _, pred = torch.max(yhat.data, 1)
    total += y.size(0)
    num_correct += (pred == y.data).sum()

print(f'Accuracy of the model on {total} test images: {num_correct * 100 / total}% ')

import matplotlib.pyplot as plt

images,labels = next(iter(testloader))

model.cpu()
output = model(images)
_,pred = torch.max(output,1)

fig = plt.figure(figsize=(10,10))
for i in range(4):
  plt.subplot(5,5,i+1)
  show_img(images[i])
  plt.title(get_class(labels[i])+"/"+get_class(pred[i]))
  plt.xticks([])
  plt.yticks([])

References:

PyTorch Autograd

Relevant Courses

July 18, 2021

python, pytorch, transfer learning