MedPipe: End-to-End Joint Search of Data Augmentation Policy 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 \textit{MedPipe}. 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 \textit{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 MedPipe can efficiently find unified networks with better
performance than previous human-designed and NAS-searched networks.