[论文解读] Teachers Do More Than Teach: Compressing Image-to-Image Models

Qing Jin, Jian Ren+5 more|arXiv (Cornell University)|Mar 5, 2021

Generative Adversarial Networks and Image Synthesis参考文献 79被引用 7

一句话总结

该论文提出CAT（Compression And Teaching）方法，利用单一教师模型联合执行架构搜索与知识蒸馏，以压缩图像到图像的GAN模型。通过引入基于Inception的残差模块和一步剪枝策略，该方法在性能与效率之间实现了最先进（SOTA）的权衡，压缩后的模型在将MACs（乘加操作数）降低数个数量级的同时，保持或超越了原始模型的FID和mIoU得分。

ABSTRACT

Generative Adversarial Networks (GANs) have achieved huge success in generating high-fidelity images, however, they suffer from low efficiency due to tremendous computational cost and bulky memory usage. Recent efforts on compression GANs show noticeable progress in obtaining smaller generators by sacrificing image quality or involving a time-consuming searching process. In this work, we aim to address these issues by introducing a teacher network that provides a search space in which efficient network architectures can be found, in addition to performing knowledge distillation. First, we revisit the search space of generative models, introducing an inception-based residual block into generators. Second, to achieve target computation cost, we propose a one-step pruning algorithm that searches a student architecture from the teacher model and substantially reduces searching cost. It requires no l1 sparsity regularization and its associated hyper-parameters, simplifying the training procedure. Finally, we propose to distill knowledge through maximizing feature similarity between teacher and student via an index named Global Kernel Alignment (GKA). Our compressed networks achieve similar or even better image fidelity (FID, mIoU) than the original models with much-reduced computational cost, e.g., MACs. Code will be released at https://github.com/snap-research/CAT.

研究动机与目标

• 为解决图像到图像GAN模型在资源受限设备上计算成本高和内存占用大的问题。
• 克服现有GAN压缩方法在图像质量下降或需要昂贵搜索过程方面的局限性。
• 通过利用教师网络同时进行架构搜索与知识蒸馏，实现高效且高保真度的模型压缩。
• 通过消除对独立超网络的需求并避免使用大量超参数的正则化方法，降低架构搜索的计算成本。
• 通过直接对齐教师与学生网络之间的特征表示（使用核对齐技术），提升蒸馏质量。

提出的方法

• 引入基于Inception的残差模块（IncResBlock）作为教师生成器的核心构建模块，为高效的学生架构提供丰富的搜索空间。
• 提出一种一步剪枝算法，直接从训练好的教师网络中移除通道以满足目标MACs预算，无需迭代或基于稀疏性的搜索。
• 设计一种基于核对齐（Kernel Alignment, KA）的知识蒸馏技术，通过最大化教师与学生特征之间的相似性，无需引入额外可学习层。
• 使用同一教师模型作为搜索空间和知识蒸馏的来源，消除了对独立超网络的需求，降低了训练复杂度。
• 在剪枝过程中仅使用一个超参数（剪枝比例），简化了训练流程并减少了调参开销。
• 通过基于KA的蒸馏端到端训练学生网络，确保学生能够从教师网络中学习到高层次的语义表征。

Figure 1 : Performance comparison between our and existing GAN compression techniques [ 11 , 20 , 36 , 64 , 70 ] on CycleGAN [ 85 ] for Horse $\shortrightarrow$ Zebra dataset. Smaller MACs indicates more efficient models. Lower FID indicates models can generate more realistic images. Our method (red

实验结果

研究问题

• RQ1 一个单一的教师模型能否有效同时作为图像到图像GAN模型压缩的搜索空间和知识蒸馏源？
• RQ2 一步剪枝是否能在无需迭代训练或ℓ¹正则化的情况下实现高效的架构搜索？
• RQ3 通过核对齐实现的特征空间相似性是否能优于依赖额外投影层的传统蒸馏方法？
• RQ4 所提方法是否在性能-效率权衡上优于现有GAN压缩技术？
• RQ5 压缩后的模型是否能在降低计算成本的同时，图像保真度超越原始模型？

主要发现

• 与之前工作相比，所提CAT方法将架构搜索成本降低了至少10,000倍，单张GPU上搜索时间低于5秒。
• 在Horse→Zebra CycleGAN基准上，压缩模型的FID（18.4）低于原始模型（19.1），且MACs减少3.5倍。
• 在Map→Aerial photo任务的Pix2pix上，压缩模型生成的河流与建筑图像比原始模型更清晰，MACs减少4.2倍。
• 在GauGAN上，压缩模型的FID为12.6，优于原始模型的13.8，且MACs减少3.8倍。
• 在所有基准测试中，该方法始终在FID和mIoU上保持或超越原始模型性能，同时大幅降低计算成本。
• 视觉结果表明，压缩模型生成的图像更具真实感和细节，尤其在原始模型失效的复杂输入上表现更优。

Figure 2 : IncResBlock includes three conventional convolution blocks and three depth-wise convolution blocks (dashed border), both with kernels sizes of 1, 3, 5. Normalization layers ( \eg , BN), and ReLU, are applied between each two consecutive convolution layers. A normalization layer that can b