What is transfer learning?
Transfer learning is a technique where a deep-learning model trained on another problem (which usually has lot of data and good accuracy for that task) is slightly modified to be used on a new problem. This is an important concept as building an entirely new model might not be take a long time or there might not be enough data for the training of that particular task. The idea is the weights/parameters of the model at the start of the layers have similar functionality and assist in better performance on the new task. Usually we freeze the weights training of the hidden layers an tweak the output layer slightly to account for the change in the task.
So for example, maybe you could have the neural network learn to recognize objects like cats and then use that knowledge or use part of that knowledge to help you do a better job reading x-ray scans. This is called transfer learning. Sometimes you can start with the weights and biases of a published netowrks as a starting point.
More details about Transfer Learning can be found on Stanford's CS231 CNN course here
In this example I use a pre-trained convolutional neural network model (ResNet-18) and modify ONLY the last layers of the model to use for our case. This model is trained on millions on images with 1000 image categories.
These two major transfer learning scenarios look as follows:
-
Finetuning the convnet: Instead of random initialization, we initialize the network with a pretrained network, like the one that is trained on imagenet 1000 dataset. Rest of the training looks as usual.
-
ConvNet as fixed feature extractor: Here, we will freeze the weights for all of the network except that of the final fully connected layer. This last fully connected layer is replaced with a new one with random weights and only this layer is trained.
This tutorial was adapted from PyTorch's official documentation (Link)
import time
import os
import copy
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision
from torchvision import datasets, models, transforms
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
# Plot matplotlib plots with white background:
%config InlineBackend.print_figure_kwargs={'facecolor' : "w"}
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device
# As defined on the tutorial page
mean = np.array([0.5, 0.5, 0.5])
std = np.array([0.25, 0.25, 0.25])
data_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
'val': transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
}
data_dir = 'data/hymenoptera_data/'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
data_transforms[x])
for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
shuffle=True, num_workers=0)
for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
print(class_names)
def imshow(inp, title):
"""Imshow for Tensor."""
fig, ax = plt.subplots(1,1, figsize=(10,10))
inp = inp.numpy().transpose((1, 2, 0))
inp = std * inp + mean
inp = np.clip(inp, 0, 1)
ax.imshow(inp)
plt.title(title)
plt.show()
# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))
# Make a grid from batch
out = torchvision.utils.make_grid(inputs)
imshow(out, title=[class_names[x] for x in classes])
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
model.train() # Set model to training mode
else:
model.eval() # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
if phase == 'train':
scheduler.step() #Step in the scheduler
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects.double() / dataset_sizes[phase]
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model for which val_acc is better
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(best_model_wts)
return model
Method 1: Fine tuning of the model
Download a pretrained model, tweak the model architecture for our use-case, and (re)train on the new dataset
- Download the model (ResNet18 in this case)
- Change the output of the final layer -- in this case from 1000 output nodes to 2 since we're looking at only 'bee' and 'ant'
- Re-train the model
Other models available from PyTorch can be viewed (here)
model = torchvision.models.resnet18(pretrained=True)
print(model)
We're interested to only change the final layer named fc -- so we will look at that layers features, the attributes for each module in the model are stored as keys.
num_features = model.fc.in_features
# Define a new linear layer as per our need -- 2 classes instead of 1000 as defined in the original
model.fc = nn.Linear(num_features, len(class_names)) #Number of classes in the end
# Send model to device
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
#Scheduler to update the learning rate for the SGD
'''
After every 7 steps in the optimizer the learning rate will be multiplied by 0.1
'''
step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size = 7, gamma = 0.1)
model = train_model(model, criterion=criterion, optimizer=optimizer, scheduler=step_lr_scheduler, num_epochs=5)
model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 2) #Number of classes in the end
model.to(device)
This block is same as before
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
#Scheduler to update the learning rate for the SGD
'''
After every 7 steps in the optimizer the learning rate will be multiplied by 0.1
'''
step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size = 7, gamma = 0.1)
Model training
model = train_model(model, criterion=criterion, optimizer=optimizer, scheduler=step_lr_scheduler, num_epochs=5)
Using fine-tuning and re-training all the weights for the new network take longer and may result in lower acceracy than having the weights fixed. When we keep the weights in earlier layers fixed, we save a lot of time in the model training aand the performance of the model is also better (provided the two tasks are quite similar). Given the network is deep, optimizing the weights for this network from scratch would have been difficult and time consuming. For such a case, transfer learning seems to be a good option.