[论文解读] Lightweight Image Super-Resolution with Hierarchical and Differentiable Neural Architecture Search

Han Huang, Shen Li+4 more|arXiv (Cornell University)|May 9, 2021

Advanced Image Processing Techniques参考文献 45被引用 10

一句话总结

本文提出了一种可微分的、分层的神经架构搜索（NAS）框架，用于轻量级单图像超分辨率（SISR），通过基于知识蒸馏的搜索空间，联合优化细胞级和网络级架构。该方法在单张GPU上仅需68G Multi-Adds（×2）和18G Multi-Adds（×4）即可实现最先进的PSNR和SSIM性能，显著优于手工设计及先前基于NAS的SISR模型，在效率与准确率方面表现更优。

ABSTRACT

Single Image Super-Resolution (SISR) tasks have achieved significant performance with deep neural networks. However, the large number of parameters in CNN-based methods for SISR tasks require heavy computations. Although several efficient SISR models have been recently proposed, most are handcrafted and thus lack flexibility. In this work, we propose a novel differentiable Neural Architecture Search (NAS) approach on both the cell-level and network-level to search for lightweight SISR models. Specifically, the cell-level search space is designed based on an information distillation mechanism, focusing on the combinations of lightweight operations and aiming to build a more lightweight and accurate SR structure. The network-level search space is designed to consider the feature connections among the cells and aims to find which information flow benefits the cell most to boost the performance. Unlike the existing Reinforcement Learning (RL) or Evolutionary Algorithm (EA) based NAS methods for SISR tasks, our search pipeline is fully differentiable, and the lightweight SISR models can be efficiently searched on both the cell-level and network-level jointly on a single GPU. Experiments show that our methods can achieve state-of-the-art performance on the benchmark datasets in terms of PSNR, SSIM, and model complexity with merely 68G Multi-Adds for $ imes 2$ and 18G Multi-Adds for $ imes 4$ SR tasks. Code will be available at \url{this https URL}.

研究动机与目标

• 通过高效架构设计，解决现有基于CNN的SISR模型计算成本高和参数负担重的问题。
• 通过引入灵活、自动化的搜索框架，克服手工设计SISR模型的局限性。
• 构建联合的细胞级与网络级搜索空间，以同时优化SISR的微观架构与宏观结构。
• 实现高效、端到端可微分的NAS用于SISR，可在单张GPU上训练，避免昂贵的强化学习或进化算法搜索流程。
• 在PSNR、SSIM和模型复杂度方面，优于现有轻量级SISR方法。

提出的方法

• 基于知识蒸馏机制设计细胞级搜索空间，优先选择能增强特征表示的轻量化操作。
• 提出一种可微分的搜索策略，支持在细胞级和网络级架构选择中进行基于梯度的优化。
• 构建网络级搜索空间，以建模细胞之间的特征流连接，识别最优的信息传播路径。
• 将细胞级与网络级搜索整合为统一的可微分流程，实现在单一训练过程中联合优化架构与权重。
• 在整个搜索过程中仅使用单张GPU，与强化学习或进化算法驱动的NAS方法相比，显著降低训练成本。
• 采用类似DARTS的可微分架构搜索方法，通过单阶段训练目标，高效搜索轻量级SISR模型。

实验结果

研究问题

• RQ1 可微分的、分层的NAS框架能否联合优化细胞级与网络级架构，从而在SISR任务中实现更高的效率与准确率？
• RQ2 基于知识蒸馏的细胞级搜索空间在SISR任务中如何在降低模型复杂度的同时提升性能？
• RQ3 与强化学习或进化算法驱动的NAS相比，完全可微分的搜索流程在SISR中能将计算成本降低多少？
• RQ4 在严格的FLOPs约束下，细胞之间最优的信息流模式是什么，能显著提升SISR性能？
• RQ5 所提方法能否在保持低参数量和低推理成本的同时，实现PSNR和SSIM的最先进性能？

主要发现

• 所提方法在×2和×4超分辨率任务的基准数据集上均实现了最先进的PSNR和SSIM性能。
• 最终搜索得到的模型仅需68G Multi-Adds（×2超分）和18G Multi-Adds（×4超分），显著降低了计算复杂度。
• 可微分的搜索流程使单张GPU上实现高效架构搜索成为可能，无需依赖昂贵的强化学习或进化算法。
• 基于知识蒸馏的细胞级搜索空间通过选择轻量化且高影响力的运算，有效平衡了准确率与效率。
• 网络级搜索识别出最优的特征连接模式，在不增加模型大小的前提下显著提升了性能。
• 该方法在准确率与模型复杂度方面均优于现有手工设计及基于NAS的SISR模型，展现出更优的效率-准确率权衡。