DataMap: Dataset transferability map for medical image classification

Abstract

Deep learning (DL)-based models especially Convolutional Neural Network (CNN) models have recently achieved great success in medical image classifications. It is usually time-consuming and labor-intensive to train a practical classification model due to the requirement of large data volume. Since medical images are more difficult to acquire and label for model training, many scholars have applied transfer learning by pre-training a model on a larger dataset and then fine-tuning it on the target dataset to obtain better classification results. However, such approach, relying on individual expertise to select related datasets, is subjective and inconsistent in performance. In this paper, we propose a simple yet effective method for measuring the transferability between different datasets, and build a Dataset Map (DataMap) that can be used as a tool to find the most relevant datasets for transfer learning on target dataset. Recent studies show that the convolutional kernels in CNN models have different function roles. Therefore, we adopt the similarity between the convolution kernels to measure the transferability between datasets. Firstly, the gradient attribution is adopted to attribute the task related convolution kernels from last few convolution layers of the same pre-trained model architecture trained with different datasets. Then, the similarity between attributed convolutional kernels is calculated to denote the transferability between different datasets. Finally, we build a DataMap with 20 medical image datasets. Extensive experimental tests on 3 mainstream CNN architectures show that the proposed method can effectively measure the transferability between different datasets. With the guidance of the DataMap, the transfer learning can achieve the best performance on various training tasks, and the accuracy of the CNN classifier can be improved by 1% to 5% through pre-training.