MedPipe: End-to-End Joint Search of Data Augmentation and Neural Architecture for 3D Medical Image Classification

Abstract

3D medical image classification with deep learning methods is challenging due to two problems, the difficulty of designing the neural architecture (NA) and the lack of large-scale 3D medical datasets for pretraining. Neural architecture search (NAS) has become a popular method to design NA automatically. Inspired by NAS, recent studies have proposed automatic data augmentation (ADA) to search the data augmentation policy (DAP) to increase the size of 3D medical datasets. However, NAS or ADA can only solve one of the above two problems. It is difficult to search both DAP and NA for 3D medical images because the computational cost of 3D data augmentation operations is high, and the joint search space is too large to find the optimal solution economically. To this end, we propose an efficient end-to-end framework for the joint search of DAP and NA applicable to 3D medical images, namely Med-Pipe. Specifically, 1) we modify all 3D data augmentation operations to be differentiable and integrate them with NA seamlessly; 2) we design a novel search space that unifies DAP and NA, i.e., each possible network is a combination of the DAP and NA, called unified network ; 3) we adopt a single-path based differentiable search algorithm to search the unified network. We evaluate our framework on nine public 3D medical datasets covering various modalities and data resolutions. Experimental results show that our Med-Pipe can efficiently find unified networks with better performance than previous human-designed and NAS-searched networks.