PNG: Micro-Structured Prune-and-Grow Networks for Flexible Image Restoration

Abstract

This paper addresses one major issue of DNN-based image restoration: the difficulty of using one model to fit multiple reconstruction requirements, such as supporting different compression ratios in neural image compression (NIC) or different zooming scales in single image super-resolution (SISR). Instead of training an independent model instance for each requirement as an individual task, we develop a practical solution that uses one model instance to support multiple requirements. We propose a general multi-task learning framework based on a novel prune-and-grow (PnG) process, where each task corresponds to each of the requirements. Different from traditional multi-task networks that use fully shared or task-specific layers, we enable in-layer partial parameter sharing to obtain both common and task-specific features at various abstraction levels. This encourages adequate sharing to improve the overall multi-task performance. The parameters are shared at a micro-structured level to both maintan the task performance and reduce inference computation. The sharing structure is automatically learned, where a model instance trained for previous tasks is progressively pruned and regrown to perform more tasks. The framework is task-generic and model-structure-agnostic. Using NIC and SISR as two example applications, extensive experiments show that the multi-task PnG network can largely reduce the overall model size and inference computation, with almost no degradation of the reconstruction performance.