Searching for Controllable Image Restoration Networks

TL;DR

This paper introduces a neural architecture search framework with task-agnostic and task-specific pruning to efficiently generate multiple image restoration effects in a single inference, significantly reducing computational cost and latency.

Contribution

It proposes a novel pruning-based approach that enables controllable image restoration with shared features, reducing inference time and computational resources.

Findings

Reduced 95.7% FLOPs for 27 effects

GPU latency decreased by 73% on 4K images

Single inference produces multiple imagery effects

Abstract

Diverse user preferences over images have recently led to a great amount of interest in controlling the imagery effects for image restoration tasks. However, existing methods require separate inference through the entire network per each output, which hinders users from readily comparing multiple imagery effects due to long latency. To this end, we propose a novel framework based on a neural architecture search technique that enables efficient generation of multiple imagery effects via two stages of pruning: task-agnostic and task-specific pruning. Specifically, task-specific pruning learns to adaptively remove the irrelevant network parameters for each task, while task-agnostic pruning learns to find an efficient architecture by sharing the early layers of the network across different tasks. Since the shared layers allow for feature reuse, only a single inference of the task-agnostic layers is needed to generate multiple imagery effects from the input image. Using the proposed task-agnostic and task-specific pruning schemes together significantly reduces the FLOPs and the actual latency of inference compared to the baseline. We reduce 95.7% of the FLOPs when generating 27 imagery effects, and make the GPU latency 73.0% faster on 4K-resolution images.