ClassPruning: Speed Up Image Restoration Networks by Dynamic N:M Pruning

TL;DR

ClassPruning dynamically adjusts the complexity of image restoration networks based on the difficulty of each image, significantly reducing computational cost while maintaining high performance.

Contribution

It introduces a novel dynamic N:M pruning approach combined with difficulty prediction to adapt network capacity per image, improving efficiency.

Findings

Reduces FLOPs by approximately 40% without performance loss.

Uses a lightweight classifier for image difficulty prediction.

Employs a new training strategy with additional loss terms.

Abstract

Image restoration tasks have achieved tremendous performance improvements with the rapid advancement of deep neural networks. However, most prevalent deep learning models perform inference statically, ignoring that different images have varying restoration difficulties and lightly degraded images can be well restored by slimmer subnetworks. To this end, we propose a new solution pipeline dubbed ClassPruning that utilizes networks with different capabilities to process images with varying restoration difficulties. In particular, we use a lightweight classifier to identify the image restoration difficulty, and then the sparse subnetworks with different capabilities can be sampled based on predicted difficulty by performing dynamic N:M fine-grained structured pruning on base restoration networks. We further propose a novel training strategy along with two additional loss terms to stabilize training and improve performance. Experiments demonstrate that ClassPruning can help existing methods save approximately 40% FLOPs while maintaining performance.