TinyGAN: Distilling BigGAN for Conditional Image Generation

知识蒸馏

该论文发表在2020年ACCV。
PDF: TinyGAN: Distilling BigGAN for Conditional Image Generation
code: https://github.com/terarachang/ACCV_TinyGAN

Abstract

生成对抗网络(GANs)在生成式图像建模中有着重要的应用，但是GANs训练很不稳定，尤其是对于大规模、复杂的数据集。最近的一个工作BigGAN极大程度上提升了在ImageNet上的图像生成质量，但该方法需要一个极大的模型，因此难以部署在资源受限的设备上。为了减少模型大小，这篇文章提出了一个黑盒知识蒸馏框架来压缩GANs，并且是一个稳定且有效的训练过程。具体方法为：将BigGAN作为教师网络，然后训练一个非常小的学生网络来模仿教师网络的功能，在Inception和FID分数两个指标上取得了不错的效果，同时生成器参数量少了16倍。

Introduction

GAN基础：GANs由生成器和判别器组成，生成器网络旨在产生和目标域相同或相似的分布，判别器网络旨在区分生成的分布和真实的数据分布，两个网络交替迭代进行训练。Conditional GANs(cGANs)是一种基于某些给定条件信息来生成样本的GAN方法，cGANs的判别器需要基于给定的条件信息来进行区分两个分布。

问题：然而GANs存在一个重要的问题：训练不稳定，特别是在大规模，复杂数据集上进行训练。最新的一个工作BigGAN，通过增大model size和batch size来减轻训练的问题，同时生成高质量的图像。但也带来了新的问题，即高计算花费和内存消耗，测试时间也较长。

那么如何进行模型压缩呢？分类任务中常用的方法有：knowledge distillation，network pruning和quantization。
本文采用的是知识蒸馏的方法，将训练好的生成器如BigGAN作为教师网络，利用其输入和输出作为训练数据，从而以一种有监督的方式训练学生网络，并且不需要获取到教师网络的内部参数。

Proposed method

方法流程如图所示。black-box指的是不需要获取到教师网络的参数或者是与教师网络共享同样的结构。

该方法分为两个步骤：

Preprocessing：即利用预训练好的BigGAN模型的生成器产生相应的样本对。
Training：利用获取到的样本对进行student网络的训练。

训练策略：
学生网络的训练包含这么几部分损失：

在训练的时候，逐渐减小pixel-level loss的权重，从而使得判别器能提供更有用的指导。

Pixel-Level Distillation Loss：

即在给定相同输入的情况下减小BigGAN和TinyGAN生成结果之间的像素距离。其中，T是参数固定的教师网络（BigGAN的生成器），S是学生网络，z是从截断的正态分布p(z)中采样得到的随机变量，y是相应的类别标签。只用该像素损失会造成生成的结果非常平滑。
Adversarial Distillation Loss: 即加入判别器进行对抗训练
Feature-Level Distillation Loss：为了进一步减轻利用像素损失生成模糊图像的问题，加入了特征级别的蒸馏损失。

其中，Di是判别器网络第i层的特征向量。

Learning from Real Distribution
为了使得模型能从ImageNet中的真实图像上进行学习，作者还还用了hinge loss作为对抗损失：

其中x是从ImageNet中采样得到的真实图像。

优点
这篇论文在方法上比较常规，但是在问题引入上还是有一定的借鉴意义。

(recoded by xt)