Notebook

In [ ]:

%matplotlib inline

Transfer Learning tutorial¶

Author: Sasank Chilamkurthy <https://chsasank.github.io>_

In this tutorial, you will learn how to train your network using transfer learning. You can read more about the transfer learning at cs231n notes <http://cs231n.github.io/transfer-learning/>__

Quoting this notes,

In practice, very few people train an entire Convolutional Network from scratch (with random initialization), because it is relatively rare to have a dataset of sufficient size. Instead, it is common to pretrain a ConvNet on a very large dataset (e.g. ImageNet, which contains 1.2 million images with 1000 categories), and then use the ConvNet either as an initialization or a fixed feature extractor for the task of interest.

These two major transfer learning scenarios looks as follows:

Finetuning the convnet: Instead of random initializaion, 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.

In [ ]:

# License: BSD# Author: Sasank Chilamkurthyfrom__future__importprint_function,divisionimporttorchimporttorch.nnasnnimporttorch.optimasoptimfromtorch.autogradimportVariableimportnumpyasnpimporttorchvisionfromtorchvisionimportdatasets,models,transformsimportmatplotlib.pyplotaspltimporttimeimportcopyimportosplt.ion()# interactive mode

Load Data¶

We will use torchvision and torch.utils.data packages for loading the data.

The problem we're going to solve today is to train a model to classify ants and bees. We have about 120 training images each for ants and bees. There are 75 validation images for each class. Usually, this is a very small dataset to generalize upon, if trained from scratch. Since we are using transfer learning, we should be able to generalize reasonably well.

This dataset is a very small subset of imagenet.

.. Note :: Download the data from here <https://download.pytorch.org/tutorial/hymenoptera_data.zip>_ and extract it to the current directory.

In [ ]:

# Data augmentation and normalization for training # Just normalization for validationdata_transforms={'train':transforms.Compose([transforms.RandomSizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]),'val':transforms.Compose([transforms.Scale(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]),}data_dir='hymenoptera_data'dsets={x:datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x])forxin['train','val']}dset_loaders={x:torch.utils.data.DataLoader(dsets[x],batch_size=4,shuffle=True,num_workers=4)forxin['train','val']}dset_sizes={x:len(dsets[x])forxin['train','val']}dset_classes=dsets['train'].classesuse_gpu=torch.cuda.is_available()

Visualize a few images ^^^^^^^^^^^^^^^^^^^^^^ Let's visualize a few training images so as to understand the data augmentations.

In [ ]:

defimshow(inp,title=None):"""Imshow for Tensor."""inp=inp.numpy().transpose((1,2,0))mean=np.array([0.485,0.456,0.406])std=np.array([0.229,0.224,0.225])inp=std*inp+meanplt.imshow(inp)iftitleisnotNone:plt.title(title)plt.pause(0.001)# pause a bit so that plots are updated# Get a batch of training datainputs,classes=next(iter(dset_loaders['train']))# Make a grid from batchout=torchvision.utils.make_grid(inputs)imshow(out,title=[dset_classes[x]forxinclasses])

Training the model¶

Now, let's write a general function to train a model. Here, we will illustrate:

Scheduling the learning rate
Saving (deep copying) the best model

In the following, parameter lr_scheduler(optimizer, epoch) is a function which modifies optimizer so that the learning rate is changed according to desired schedule.

In [ ]:

deftrain_model(model,criterion,optimizer,lr_scheduler,num_epochs=25):since=time.time()best_model=modelbest_acc=0.0forepochinrange(num_epochs):print('Epoch {}/{}'.format(epoch,num_epochs-1))print('-'*10)# Each epoch has a training and validation phaseforphasein['train','val']:ifphase=='train':optimizer=lr_scheduler(optimizer,epoch)model.train(True)# Set model to training modeelse:model.train(False)# Set model to evaluate moderunning_loss=0.0running_corrects=0# Iterate over data.fordataindset_loaders[phase]:# get the inputsinputs,labels=data# wrap them in Variableifuse_gpu:inputs,labels=Variable(inputs.cuda()), \ Variable(labels.cuda())else:inputs,labels=Variable(inputs),Variable(labels)# zero the parameter gradientsoptimizer.zero_grad()# forwardoutputs=model(inputs)_,preds=torch.max(outputs.data,1)loss=criterion(outputs,labels)# backward + optimize only if in training phaseifphase=='train':loss.backward()optimizer.step()# statisticsrunning_loss+=loss.data[0]running_corrects+=torch.sum(preds==labels.data)epoch_loss=running_loss/dset_sizes[phase]epoch_acc=running_corrects/dset_sizes[phase]print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase,epoch_loss,epoch_acc))# deep copy the modelifphase=='val'andepoch_acc>best_acc:best_acc=epoch_accbest_model=copy.deepcopy(model)print()time_elapsed=time.time()-sinceprint('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60,time_elapsed%60))print('Best val Acc: {:4f}'.format(best_acc))returnbest_model

Learning rate scheduler ^^^^^^^^^^^^^^^^^^^^^^^ Let's create our learning rate scheduler. We will exponentially decrease the learning rate once every few epochs.

In [ ]:

defexp_lr_scheduler(optimizer,epoch,init_lr=0.001,lr_decay_epoch=7):"""Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""lr=init_lr*(0.1**(epoch//lr_decay_epoch))ifepoch%lr_decay_epoch==0:print('LR is set to {}'.format(lr))forparam_groupinoptimizer.param_groups:param_group['lr']=lrreturnoptimizer

Visualizing the model predictions ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Generic function to display predictions for a few images

In [ ]:

defvisualize_model(model,num_images=6):images_so_far=0fig=plt.figure()fori,datainenumerate(dset_loaders['val']):inputs,labels=dataifuse_gpu:inputs,labels=Variable(inputs.cuda()),Variable(labels.cuda())else:inputs,labels=Variable(inputs),Variable(labels)outputs=model(inputs)_,preds=torch.max(outputs.data,1)forjinrange(inputs.size()[0]):images_so_far+=1ax=plt.subplot(num_images//2,2,images_so_far)ax.axis('off')ax.set_title('predicted: {}'.format(dset_classes[preds[j]]))imshow(inputs.cpu().data[j])ifimages_so_far==num_images:return

Finetuning the convnet¶

Load a pretrained model and reset final fully connected layer.

In [ ]:

model_ft=models.resnet18(pretrained=True)num_ftrs=model_ft.fc.in_featuresmodel_ft.fc=nn.Linear(num_ftrs,2)ifuse_gpu:model_ft=model_ft.cuda()criterion=nn.CrossEntropyLoss()# Observe that all parameters are being optimizedoptimizer_ft=optim.SGD(model_ft.parameters(),lr=0.001,momentum=0.9)

Train and evaluate ^^^^^^^^^^^^^^^^^^

It should take around 15-25 min on CPU. On GPU though, it takes less than a minute.

In [ ]:

model_ft=train_model(model_ft,criterion,optimizer_ft,exp_lr_scheduler,num_epochs=25)

In [ ]:

visualize_model(model_ft)

ConvNet as fixed feature extractor¶

Here, we need to freeze all the network except the final layer. We need to set requires_grad == False to freeze the parameters so that the gradients are not computed in backward().

You can read more about this in the documentation here <http://pytorch.org/docs/notes/autograd.html#excluding-subgraphs-from-backward>__.

In [ ]:

model_conv=torchvision.models.resnet18(pretrained=True)forparaminmodel_conv.parameters():param.requires_grad=False# Parameters of newly constructed modules have requires_grad=True by defaultnum_ftrs=model_conv.fc.in_featuresmodel_conv.fc=nn.Linear(num_ftrs,2)ifuse_gpu:model_conv=model_conv.cuda()criterion=nn.CrossEntropyLoss()# Observe that only parameters of final layer are being optimized as# opoosed to before.optimizer_conv=optim.SGD(model_conv.fc.parameters(),lr=0.001,momentum=0.9)

Train and evaluate ^^^^^^^^^^^^^^^^^^

On CPU this will take about half the time compared to previous scenario. This is expected as gradients don't need to be computed for most of the network. However, forward does need to be computed.

In [ ]:

model_conv=train_model(model_conv,criterion,optimizer_conv,exp_lr_scheduler,num_epochs=25)

In [ ]:

visualize_model(model_conv)plt.ioff()plt.show()