Return to search

Pruning of U-Nets : For Faster and Smaller Machine Learning Models in Medical Image Segmentation

Accurate medical image segmentation is crucial for safely and effectively administering radiation therapy in cancer treatment. State of the art methods for automatic segmentation of 3D images are currently based on the U-net machine learning architecture. The current U-net models are large, often containing millions of parameters. However, the size of these machine learning models can be decreased by removing parts of the models, in what is called pruning. One algorithm, called simultaneous training and pruning (STAMP) has shown capable of reducing the model sizes upwards of 80% while keeping similar or higher levels of performance for medical image segmentation tasks.  This thesis investigates the impact of using the STAMP algorithm to reduce model size and inference time for medical image segmentation on 3D images, using one MRI and two CT datasets. Surprisingly, we show that pruning convolutional filters randomly achieves performance comparable, if not better than STAMP, provided that the filters are always removed from the largest parts of the U-net.  Inspired by these results, a modified "Flat U-net" is proposed, where an equal number of convolutional filters are used in all parts of the U-net, similar to what was achieved after pruning with our simplified pruning algorithm. The modified U-net achieves similar levels of test dice score as both a regular U-net and the STAMP pruning algorithm, on multiple datasets while avoiding pruning altogether. In addition to this the proposed modification reduces the model size by more than a factor of 12, and the number of computations by around 35%, compared to a normal U-net with the same number of input-layer convolutional filters.

Identiferoai:union.ndltd.org:UPSALLA1/oai:DiVA.org:uu-531810
Date January 2024
CreatorsHassler, Ture
PublisherUppsala universitet, Avdelningen Vi3
Source SetsDiVA Archive at Upsalla University
LanguageEnglish
Detected LanguageEnglish
TypeStudent thesis, info:eu-repo/semantics/bachelorThesis, text
Formatapplication/pdf
Rightsinfo:eu-repo/semantics/openAccess
RelationUPTEC F, 1401-5757 ; 24011

Page generated in 0.0019 seconds