Efficient Model-Based Deep Learning via Network Pruning and Fine-Tuning

TL;DR

This paper introduces a structured pruning and fine-tuning approach for model-based deep learning networks, significantly reducing computational complexity while maintaining performance in imaging inverse problems.

Contribution

It presents a novel method to prune and fine-tune MBDL networks, enhancing their scalability and efficiency in large-scale imaging applications.

Findings

Pruning reduces parameters in MBDL networks by up to 50%.

Fine-tuning mitigates performance loss after pruning.

Achieves 50% and 32% acceleration in DEQ and DU methods with minimal accuracy impact.

Abstract

Model-based deep learning (MBDL) is a powerful methodology for designing deep models to solve imaging inverse problems. MBDL networks can be seen as iterative algorithms that estimate the desired image using a physical measurement model and a learned image prior specified using a convolutional neural net (CNNs). The iterative nature of MBDL networks increases the test-time computational complexity, which limits their applicability in certain large-scale applications. Here we make two contributions to address this issue: First, we show how structured pruning can be adopted to reduce the number of parameters in MBDL networks. Second, we present three methods to fine-tune the pruned MBDL networks to mitigate potential performance loss. Each fine-tuning strategy has a unique benefit that depends on the presence of a pre-trained model and a high-quality ground truth. We show that our pruning and fine-tuning approach can accelerate image reconstruction using popular deep equilibrium learning (DEQ) and deep unfolding (DU) methods by 50% and 32%, respectively, with nearly no performance loss. This work thus offers a step forward for solving inverse problems by showing the potential of pruning to improve the scalability of MBDL. Code is available at https://github.com/wustl-cig/MBDL_Pruning .