Sauron U-Net: Simple automated redundancy elimination in medical image segmentation via filter pruning

원문 초록

We introduce Sauron, a filter pruning method that eliminates redundant feature maps of convolutional neural networks (CNNs). Sauron optimizes, jointly with the loss function, a regularization term that promotes feature maps clustering at each convolutional layer by reducing the distance between feature maps. Sauron then eliminates the filters corresponding to the redundant feature maps by using automatically adjusted layer-specific thresholds. Unlike most filter pruning methods, Sauron requires minimal changes to typical neural network optimization because it prunes and optimizes CNNs jointly, which, in turn, accelerates the optimization over time. Moreover, unlike with other cluster-based approaches, the user does not need to specify the number of clusters in advance, a hyperparameter that is difficult to tune. We evaluated Sauron and five state-of-the-art filter pruning methods on four medical image segmentation tasks. This is an area where little attention has been paid to filter pruning, but where smaller CNN models are desirable for local deployment, mitigating privacy concerns associated with cloud-based solutions. Sauron was the only method that achieved a reduction in model size of over 90% without deteriorating substantially the performance. Sauron also achieved, overall, the fastest models at inference time in machines with and without GPUs. Finally, we show through experiments that the feature maps of models pruned with Sauron are highly interpretable, which is essential for medical image segmentation.