STAMP: Simultaneous Training and Model Pruning for Low Data Regimes in Medical Image Segmentation

Medical Image Analysis 2022
Nicola Dinsdale1,2
Mark Jenkinson1,3,4
Ana Namburete2

1Wellcome Centre for Integrative Neuroimaging, University of Oxford
2OMNI, Department of Computer Science, University of Oxford
3Australian Institute for Machine Learning (AIML), University of Adelaide
4South Australian Health and Medical Research Institute (SAHMRI)

[Paper]
[Code]

Abstract

Acquisition of high quality manual annotations is vital for the development of segmentation algorithms. However, to create them we require a substantial amount of expert time and knowledge. Large numbers of labels are required to train convolutional neural networks due to the vast number of parameters that must be learned in the optimisation process. Here, we develop the STAMP algorithm to allow the simultaneous training and pruning of a UNet architecture for medical image segmentation with targeted channelwise dropout to make the network robust to the pruning. We demonstrate the technique across segmentation tasks and imaging modalities. It is then shown that, through online pruning, we are able to train networks to have much higher performance than the equivalent standard UNet models while reducing their size by more than 85% in terms of parameters. This has the potential to allow networks to be directly trained on datasets where very low numbers of labels are available.

Pipeline

Results

Comparison to Standard UNet

We simultaneously train and prune UNet architectures, using the STAMP algorithm and Targeted Dropout. We apply the algorithm to a range of modalities and segmentation tasks.
We show that we can reduce the number of parameters in the model to a fraction of the orignal size and improve the performance compared to the standard UNet training.


Performance across tasks and modalities

We show that the results are consistent across tasks and modalities, with STAMP+ out performing the Standard UNet approach.


Performance in low data regimes

We also show that STAMP+ significantly improves performance in very low data, low label regimes.

Acknowledgements

ND is supported by the Engineering and Physical Sciences Research Council (EPSRC) and Medical Research Council (MRC) [grant number EP/L016052/1]. MJ is supported by the National Institute for Health Research (NIHR), Oxford Biomedical Research Centre (BRC), and this research was funded by the Wellcome Trust [215573/Z/19/Z]. The Wellcome Centre for Integrative Neuroimaging is supported by core funding from the Wellcome Trust [203139/Z/16/Z]. AN is grateful for support from the UK Royal Academy of Engineering under the Engineering for Development Research Fellowships scheme. The computational aspects of this research were supported by the Wellcome Trust Core Award [Grant Number 203141/Z/16/Z] and the NIHR Oxford BRC. The views expressed are those of the author(s) and not necessarily those of the NHS, the NIHR or the Department of Health. This template was originally made by Phillip Isola and Richard Zhang for a colorful ECCV project; the code can be found here.