AutoTune: Automatically Tuning Convolutional Neural Networks for Improved Transfer Learning

Highlights

• This paper introduces a framework called AutoTune which finds the optimal number of layers to be fine-tuned automatically for a target dataset for improved transfer learning.

• Bayesian Optimization technique is applied to learn the pre-trained CNN layers with the knowledge of the target dataset.

• Several experiments are performed on CalTech-101, CalTech-256, and Stanford Dogs using the proposed AutoTune to establish the efficacy of the proposed model compared to state-of-the-art.

• In order to validate the robustness of the proposed framework across different CNN architectures, we experiment with three popular CNNs, namely, VGG-16 (Simonyan and Zisserman, 2014), ResNet-50 (He et al., 2016), and DenseNet-121 (Huang et al., 2017).

Abstract

Transfer learning enables solving a specific task having limited data by using the pre-trained deep networks trained on large-scale datasets. Typically, while transferring the learned knowledge from source task to the target task, the last few layers are fine-tuned (re-trained) over the target dataset. However, these layers are originally designed for the source task that might not be suitable for the target task. In this paper, we introduce a mechanism for automatically tuning the Convolutional Neural Networks (CNN) for improved transfer learning. The pre-trained CNN layers are tuned with the knowledge from target data using Bayesian Optimization. First, we train the final layer of the base CNN model by replacing the number of neurons in the softmax layer with the number of classes involved in the target task. Next, the CNN is tuned automatically by observing the classification performance on the validation data (greedy criteria). To evaluate the performance of the proposed method, experiments are conducted on three benchmark datasets, e.g., CalTech-101, CalTech-256, and Stanford Dogs. The classification results obtained through the proposed AutoTune method outperforms the standard baseline transfer learning methods over the three datasets by achieving 95.92%, 86.54%, and 84.67% accuracy over CalTech-101, CalTech-256, and Stanford Dogs, respectively. The experimental results obtained in this study depict that tuning of the pre-trained CNN layers with the knowledge from the target dataset confesses better transfer learning ability. The source codes are available at https://github.com/JekyllAndHyde8999/AutoTune_CNN_TransferLearning.

Introduction

The ability of Convolutional Neural Networks (CNN) to perform feature extraction and decision making in one-shot creates enormous demand in several application areas, such as object recognition (Krizhevsky, Sutskever, & Hinton, 2012), language translation (Zhang, Zong, et al., 2015), and many more. However, the performance of the deep learning models is sensitive w.r.t. the small changes made in both network hyperparameters settings, such as the number of layers, filter dimension of a convolution layer, etc. and it is also sensitive to the other training parameters, such as learning rate, activation function, and so on. Most of the CNNs available in the literature are carefully designed in terms of these hyperparameters by the domain experts (He et al., 2016, Krizhevsky et al., 2012, Simonyan and Zisserman, 2014, Szegedy et al., 2015).

In recent years, researchers have made substantial efforts to automatically learning the structure of a CNN for a specific task (Liu et al., 2018, Zoph et al., 2018), known as Neural Architecture Search (NAS). Although these methods out-perform most of the hand-engineered architectures, the search process requires huge computational resources and time to train the proxy CNNs that are explored during the architecture search process (Zoph & Le, 2016). Particularly while working on small datasets, the process of NAS becomes a challenge. Transfer learning is a popularly adopted technique to reduce the demand for both large computational resources and training data by providing promising performance over small datasets.

Many machine learning algorithms assume that the training (source) data and future data (target) have the similar distribution. In this direction, many Metric-learning algorithms (Dong et al., 2017, Hu et al., 2015) are proposed in the literature. To mention a few, Deep Transfer Metric Learning (DTML) method introduced by Hu et al. (2015), in which a discriminative distance network is trained for cross-dataset visual recognition by maximizing the inter-class variations and minimizing the intra-class similarity. Similarly, Dong et al. (2017) proposed an Ensemble Discriminative Local Metric Learning (EDLML) which aims at learning a sub-space to keep all the intra-class samples as close as possible, while the samples belong to different classes are well separated. Shi, Du, and Zhang (2015) presented a semi-supervised domain adaptation approach which finds new representations of the images belong to the source domain using multiple linear transformations.

Typically, while transforming the learned knowledge from source task to the target task, the classification layer of the pre-trained CNN is dropped, after which a new softmax layer is stacked that is trained over the target dataset during transfer learning. The number of layers to be fine-tuned can be decided majorly based on the size of the target dataset and the similarity between the source and target datasets (Karpathy & Johnson, 2017). However, the pre-trained CNN model is designed for the source dataset, which may not perform well over the target dataset.

In this paper, we attempt to automatically tune the pre-trained CNN to make it suitable for the target task/dataset. To achieve this objective, initially, we drop the softmax layer of the pre-trained CNN by replacing it with a new softmax layer having the neurons equal to the number of classes in the target dataset. Next, we automatically tune the layers of CNN using Bayesian Optimization (Frazier, 2018). It is a well-known fact that the initial layers of CNN represent primitive features such as edges and blobs which are generic to many tasks. On the other hand, the final layers of CNN represent the features that are very specific to the learning task (Zeiler & Fergus, 2014). Based on the above idea, in our work, the layers of the CNN are tuned from right to left (i.e., from final layers to initial layers) by observing the network performance on the validation data. The results obtained through our experiments indicate that tuning the optimum number of layers with respect to the target task leads to better image classification performance in the context of transfer learning. Fig. 1 shows an overview of the proposed idea of improving the transfer learning process. Next, we provide a survey of literature in the specific research area focused in this study.