GAN Slimming: All-in-One GAN Compression by A Unified Optimization Framework

Introduction

Generative adversarial networks (GANs) , especially, image-to-image translation GANs, have been successfully applied to image synthesis , style transfer , image editing and enhancement , to name just a few. Due to the growing usage, there has been an increasing demand to deploy them on resource-constrained devices . For example, many filter-based image editing applications now desire to run image-to-image translation GANs locally. However, GANs, just like most other deep learning models, bear explosive parameter amounts and computational complexities. For example, in order to process a $256\times 256$ image, a state-of-the-art style transfer network, CartoonGAN , would cost over 56 GFLOPs. Launching such models on mobile devices requires considerable memory and computation costs, which would be infeasible for most devices, or at least degrades user experience due to the significant latency.

Existing deep model compression methods mainly focus on image classification or segmentation tasks, and were not directly applicable on GAN compression tasks due to notorious instability of GAN minimax training. For example, shows that generators compressed by state-of-the-art classifier compression methods all suffer great performance decay compared with the original generator. Combining (either heuristically cascading or jointly training) multiple different compression techniques, such as channel pruning, model distillation, quantization and weight sharing, has been shown to outperform separately using single compression techniques alone in traditional classification tasks . In comparison, current methods have so far only tried to apply one single technique to compressing GANs. proposed the first dedicated GAN compression algorithm: an evolutionary method based channel pruning algorithm. However, the method is specifically designed for CycleGAN and non-straightforward to extend to GANs without cycle consistency structure (e.g., encoder-decoder GANs that are also popular). A latest work proposed to training an efficient generator by model distillation. By removing the dependency on cycle consistency structure, achieves more general-purpose GAN compression than . However, the student network in is still hand-crafted and relies on significant architecture engineering for good performance.

As discussed in , applying a single compression technique to GANs is already challenging due to their notorious training instability. As one may imagine, integrating multiple compression techniques together for GAN compression will only further amplify the instability, putting an open question:

Can we gain more from combining multiple compression means for GANs? If yes, how to overcome the arising challenge of GAN instability?

Our answer is by presenting the first end-to-end optimization framework combining multiple compression means for general GAN compression, named GAN Slimming (GS). The core contribution of GS is a unified optimization form, that seamlessly integrates three popular model compression techniques (channel pruning, quantization and model distillation), to be jointly optimized in a minimax optimization framework. GS pioneers to advance GAN compression into jointly leveraging multiple model compression methods, and demonstrate the feasibility and promise of doing so, despite the GAN instability.

Experiments demonstrate that GS overwhelms state-of-the-art GAN compression options that rely on single compression means. For example, we compress the heavily-parameterized CartoonGAN by up to $47\times$, achieving nearly real-time cartoon style transfer on mobile devices, with minimal visual quality loss. Moreover, we have included a detailed ablation study for a deeper understanding of GS. Specifically, we demonstrate that naively stacking different compression techniques cannot achieve satisfactory GAN compression, sometimes even hurting catastrophically, therefore testifying the necessity of our unified optimization. We also verify the effectiveness of incorporating the minimax objective into this specific problem.

Related Works

Quantization, as another popular compression means, reduces the bit width of the element-level numerical representations of weights and activations. Earlier works presented to quantize all layer-wise weights and activations to the same low bit width, e.g., from 32 bits to 8 bits or less. The model could even consist of only binary weights in the extreme case . Note that, introducing quantization into network weights or activations will result in notable difficulty for propagating gradients . Straight-through estimator (STE) is a successful tool to solve this problem by a proxy gradient for back propagation.

Knowledge distillation was first developed in to transfer the knowledge in an ensemble of models to a single model, using a soft target distribution produced by the former models. It was later on widely used to obtain a smaller network (student model), by fitting the “soft labels” (probabilistic outputs) generated from a trained larger network (teacher model). used distillation to train a more efficient and accurate predictor. unified distillation and privileged information into one generalized distillation framework to learn better representations. used generative adversarial training for model distillation.

For compressing a deep classifier, proposed to jointly train (unstructured) pruning and quantization together. adopted knowledge distillation to fine-tune a pruned student network, by utilizing the original dense network as teacher, which essentially followed a two-step cascade pipeline. Similarly, used full-precision networks as teachers to distill low-precision student networks. showed jointly training pruning and quantization can obtain compact classifiers with state-of-the-art trade-off between model efficiency and adversarial robustness.

Up to our best knowledge, all above methods cascade or unify two compression techniques, besides that they investigate compressing deep classifiers only. In comparison, our proposed framework jointly optimize three methods in one unified formA concurrent work jointly optimized pruning, decomposition, and quantization, into one unified framework for reducing the memory storage/access., that is innovative even for general model compression. It is further adapted for the special GAN scenario, by incorporating the minimax loss.

2 GAN Compression

GANs have been successful on many image generation and translation tasks , yet their training remains notoriously unstable. Numerous techniques were developed to stabilize the GAN training, e.g., spectral normalization , gradient penalty and progressive training . As discussed in , the training difficulty causes extra challenges for compressing GANs, and failed many traditional pruning methods for classifiers such as .

The authors of proposed the first dedicated GAN compression method: a co-evolution algorithm based channel pruning method for CycleGAN. Albeit successfully demonstrated on the style transfer application, their method faces several limitations. First, their co-evolution algorithm relies on the cycle consistency loss to simultaneously compress generators of both directions. It is hence non-straightforward to extend to image-to-image GANs without cycle consistent loss (e.g., encoder-decoder GANs ). Second, in order to avoid the instability in GAN training, the authors model GAN compression as a “dense prediction” process by fixing the original discriminator instead of jointly updating it with the generator in a minimax optimization framework. This surrogate leads to degraded performance of the compressed generator, since the fixed discriminator may not suit the changed (compressed) generator capacity. These limitations hurdle both its broader application scope and performance.

The latest concurrent work explored model distillation: to guide the student to effectively inherit knowledge from the teacher, the authors proposed to jointly distill generator and discriminator in a minimax two-player game. improved over by removing the above two mentioned hurdles. However, as we observe from experiments (and also confirmed with their authors), the success of hinges notably on the appropriate design of student network architectures. Our method could be considered as another important step over , that “learns” the student architecture jointly with the distillation, via pruning and quantization, as to be explained by the end of Section 3.1.

The GAN Slimming Framework

Considering a dense full-precision generator $G_{0}$ which converts the images from one domain ${\mathcal{X}}$ to another ${\mathcal{Y}}$, our aim is to obtain a more efficient generator $G$ from $G_{0}$, such that their generated images $\{G_{0}(x),x\in{\mathcal{X}}\}$ and $\{G(x),x\in{\mathcal{X}}\}$ have similar style transfer qualities. In this section, we first outline the unified optimization form of our GS framework combining model distillation, channel pruning and quantization (Section 3.1). We then show how to solve each part of the optimization problem respectively (Section 3.2), and eventually present the overall algorithm (Section 3.3).

We start formulating our GS objective from the traditional minimax optimization problem in GAN:

where $D$ is the discriminator jointly trained with efficient generator $G$ by minimax optimization. Since $G$ is the functional part to be deployed on mobile devices and $D$ can be discarded after training, we do not need to compress $D$. Inspired by the success of model distillation in previous works , we add a model distillation loss term $L_{dist}$ to enforce the small generator $G$ to mimic the behaviour of original large generator $G_{0}$, where $\textrm{d}(\cdot,\cdot)$ is some distance metric:

The remaining key question is: how to properly define the architecture of $G$? Previous methods first hand-crafted the smaller student model’s architecture and then performed distillation. However, it is well known that the choice of the student network structure will affect the final performance notably too, in addition to the teacher model’s strength.

Unlike existing distillation methods , we propose to jointly infer the $G$ architecture together with the distillation process. Specifically, we assume that $G$ can be “slimmed” from $G_{0}$, through two popular compression operations: channel pruning and quantization. For channel pruning, we follow to apply $L_{1}$ norm on the trainable scale parameters $\gamma$ in the normalization layers to encourage channel sparsity: $L_{cp}=\|\gamma\|_{1}$. Denoting all other trainable weights in $G$ as $W$, we could incorporate the channel pruning via such sparsity constraint into the distillation loss in Eq. (2) as below:

where $\rho$ is the trade-off parameter controlling the network sparsity level. Further, to integrate quantization, we propose to quantize both activations and weights,We only quantize $W$, while always leaving $\gamma$ unquantized. using two quantizers $q_{a}(\cdot)$ and $q_{w}(\cdot)$, respectively, to enable the potential flexibility for hybrid quantization . While it is completely feasible to adopt learnable quantization intervals , we adopt uniform quantizers with pre-defined bit-width for $q_{a}(\cdot)$ and $q_{w}(\cdot)$, respectively, for the sake of simplicity (including hardware implementation ease). The quantized weights can be expressed as $q_{w}(W)$, while we use $G_{q}$ to denote generators equipped with activation quantization $q_{a}(\cdot)$ for notation compactness. Eventually, the final objective combining model distillation, channel pruning and quantization has the following form:

where $\theta$ represents the parameters in $D$. The blue parts represent the distillation component, green represents channel pruning red represents quantization. The above Eq. (4) is the target objective of GS, which is to be solved in a minimax optimization framework:

Our framework could be alternatively interpreted as performing a special neural architecture search (NAS) to obtain the student model, where the student’s architecture needs be “morphable” from the teacher’s through only pruning and quantization operations.

Interestingly, two concurrent works have successfully applied NAS to search efficient generator architectures, and both achieved very promising performance too. We notice that a notable portion of the performance gains shall be attributed to the carefully designed search spaces, as well as computationally intensive search algorithms. In comparison, our framework is based on an end-to-end optimization formulation, that (1) has explainable and well-understood behaviors; (2) is lighter and more stable to solve; and (3) is also free of the NAS algorithm’s typical engineering overhead (such as defining the search space and tuning search algorithms). Since our method directly shrinks the original dense model via pruning and quantization only, it cannot introduce any new operator not existing in the original model. That inspires us to combine the two streams of compression ideas (optimization-based versus NAS-based), as future work.

2 End-to-End Optimization

The difficulties of optimizing (5) can be summarized in three-folds. First, the minimax optimization problem itself is unstable. Second, updating $W$ involves non-differentiable quantization operations. Third, updating $\gamma$ involves a sparse loss term that is also non-differentiable. Below we discuss how to optimize them.

The sub-problem for updating $W$ in Eq. (5) is:

To solve (LABEL:eq:W_step) with gradient-based methods, we need to calculate $\nabla_{W}L_{W}$, which is difficult due to the non-differentiable $q_{a}(\cdot)$ and $q_{w}(\cdot)$. We now define the concrete form of $q_{a}(\cdot)$, $q_{w}(\cdot)$ and then demonstrate how to back propagate through them in order to calculate $\nabla_{W}L_{W}$. Since both $q_{a}(\cdot)$ and $q_{w}(\cdot)$ are elementwise operations, we only discuss how they work on scalars. We use $a$ and $w$ to denote a scalar element in the activation and convolution kernel tensors respectively.

When quantizing activations, we first clamp activations into range $[0,p]$ to bound the values, and then use $s_{a}={p/2^{m}}$ as a scale factor to convert the floating point number to $m$ bits integers: $\text{round}(\min(\max(0,a),p)/s_{a})$. Thus the activation quantization operator is as follows:

For weights quantization, we keep the range of the original weights and use symmetric coding for positive and negative ranges to quantize weights to $n$ bits. Specifically, the scale factor $s_{w}={\|w\|_{\infty}/2^{(n-1)}}$, leading to the quantization operator for weights:

Since both quantization operators are non-differentiable, we use a proxy as the “pseudo” gradient in the backward pass, known as the straight through estimator (STE). For the activation quantization, we use

Similarly for the weight quantization, the pseudo gradient is set to

Now that we have defined the derivatives of $q_{a}(\cdot)$ and $q_{w}(\cdot)$, we can calculate $\nabla_{W}L_{W}$ through back propagation and update $W$ using the Adam optimizer .

2.2 Updating γ𝛾\gamma

The sub-problem for updating $\gamma$ in Eq. (4) is a sparse optimization problem with a non-conventional fidelity term:

We use the proximal gradient to update $\gamma$ as follows:

where $\odot$ is element-wise product, $\textrm{sgn}(\cdot)$ and $\textrm{max}(\cdot,\cdot)$ are element-wise sign and maximum functions respectively.

2.3 Updating θ𝜃\theta

The sub-problem of updating $\theta$ is the inner maximization problem in Eq. (5), which we solve by the gradient ascent method:

We iteratively update $D$ (parameterized by $\theta$) and $G$ (parameterized by $W$ and $\gamma$) following .

3 Algorithm Implementation

Equipped with the above gradient computation, the last missing piece in solving problem (4) is to choose $d$. Note that, most previous distillation works are for classification-type models with softmax outputs (soft labels), and therefore adopt KL divergence. For GAN compression, the goal of distillation shall minimize the discrepancy between two sets of generated images. To this end, we adopt the perceptual loss as our choice of $d$. It has shown to effectively measure not only low-level visual cue, but also high-level semantic differences between images, and has been popularly adopted to regularizing GAN-based image generation.

Finally, Algorithm 1 summarizes our GS algorithm with end-to-end optimization. By default, we quantize both activation and kernel weights uniformly to 8-bit (i.e., $m$ = $n$ = $8$) and set activation clamping threshold $p$ to $4$. We use Adam ($\beta_{1}$ = $0.9$, $\beta_{2}$ = $0.5$, following ) to update $W$ and $\theta$, and SGD to update $\gamma$. We also use two groups of learning rates $\alpha^{(t)}$ and $\eta^{(t)}$ for updating $\{W,\theta\}$ and $\gamma$ respectively. $\alpha^{(t)}$ starts to be decayed linearly to zero, from the $T/2$-th iteration, while $\eta^{(t)}$ is decayed using a cosine annealing scheduler.

Experiments

Image translation and stylization is currently an important motivating application to deploy GANs on mobile devices. In this section, we compare GS with the only two published GAN compression methods CEC and GD on horse2zebra and summer2winter datasets. Following , we use model size and FLOPs to measure the efficiency of generator and use FID between source style test set transfer results and target style test set to quantitatively measure the effectiveness of style transfer. We used the same implementation of FID as for fair comparison. The metric statistic of original CycleGAN is summarized in Table 1. We denote the original dense model as $G_{0}$ and an arbitrary compressed generator as $G$. Following , we further define the following three metrics to evaluate efficiency-quality trade-off of different compression methods:

Larger $r_{c}$ and $r_{s}$ indicate more model compactness and efficiency and larger $r_{f}$ indicates better style transfer quality.

Quantitative comparison results on four different tasks are shown in Table 2. GS-32 outperforms both CEC and GD on all four tasks, in terms that it achieves better FID (larger $r_{f}$) with less computational budgets (larger $r_{c}$ and $r_{s}$). For example, on horse-to-zebra task, GS-32 has much better FID than both CEC and GD, while achieving more model compactness. Combined with quantization, our method can further boost the model efficiency (much larger $r_{c}$) with minimal loss of performance (similar $r_{f}$). For example, on horse-to-zebra task, GS-8 achieves $4\times$ larger $r_{c}$ compared with GS-32 with negligible FID drop. On winter-to-summer task, GS-8 compress CycleGAN by $31\times$ and achieve even slightly better FID. The visual comparison results are collectively displayed in Fig. 2. We compare the transfer results of four images reported in for fair comparison. As we can see, the visual quality of GS is better than or at least comparable to those of CEC and GD.

2 Ablation study

In order to show the superiority of our unified optimization framework over single compression methods and their naive combinations, we conduct thorough ablation studies by comparing the following methods:

Distillation (i.e., GD ): Use model distillation alone to train a slim student generator.Following , we use student networks with $1/2$ channels of the original generator.

Channel pruning (CP): Directly use channel pruning during GAN minimax training process. This is implemented by adding $L_{cp}$ to $L_{GAN}$. After channel pruning, we finetune the sub-network by minimax optimizing Eq. (1).

Cascade: Distillation + CP (D+CP): Use channel pruning to further compress on the student network obtained by model distillation. Then finetune the sub-network by minimax optimizing Eq. (1).

Cascade: CP + Distillation (CP+D): First do channel pruning on the original network, then use distillation to finetune the pruned network. This method is shown to outperform using channel pruning alone on classification tasks .

GS-32: Jointly optimizing channel pruning and distillation.

Cascade: GS-32 + quantization (postQ): First use GS-32 to compress the original network, then use 8 bit quantization as post processing and also do quantization-aware finetune on the quantized model by solving problem (1).

GS-8: Jointly optimizing channel pruning, distillation and quantization.

GS-8 (MSE): Replace the perceptual loss in GS-8 by MSE loss.

GAN compression with fixed discriminator (i.e., CEC ): Co-evolution based channel pruning. Modeling GAN compression as dense prediction process instead of minimax problem by fixing the discriminator (both network structure and parameter values) during compression process.

Numerical and visualization results on horse2zebra dataset are shown in Fig. 3 and Fig. 4 respectively. As we can see, our unified optimization method achieves superior trade-off between style transfer quality and model efficiency compared with single compression techniques used separately (e.g., CP, GD) and their naive combinations (e.g., CP+D, D+CP, postQ), showing the effectiveness of our unified optimization framework. For example, directly injecting channel sparsity in GAN minimax optimization (CP) greatly increases the training instability and achieves degraded image generation quality as shown in Fig. 4. This aligns with the conclusions in that model compression methods developed for classifiers are not directly applicable on GAN compression tasks. Using model distillation to finetune channel pruned models (CP+D) can indeed improve image generation quality compared with CP, however the generation quality is still much more inferior to our methods at similar compression ratio, as shown in Fig. 3 and Fig. 4. Compared with GD, which uses a hand-crafted student network, GS-32 achieves much better FID with even considerably larger compression ratio, showing the effectiveness of jointly searching slim student network structures by channel pruning and training the student network with model distillation. In contrast, directly using channel pruning to further compress the student generator trained by GD (D+CP) will catastrophically hurt the image translation performance. Doing post quantization and quantization-aware finetune (postQ) on GS-32 models also suffers great degradation in image translation quality compared with GS-8, showing the necessity to jointly train quantization with channel pruning and model distillation in our unified optimization framework. Replacing perceptual loss with MSE loss as $d$ in Eq. (2) fails to generate satisfying target images, since MSE loss cannot effectively capture the high-level semantic differences between images. Last but not least, GS-32 largely outperforms CEC, verifying the effectiveness of incorporating minimax objective into GAN compression problem.

3 Real-world Application: CartoonGAN

Finally, we apply GS to a recently proposed style transfer network CartoonGAN, which transforms photos to cartoon images, in order to deploy the model on mobile devices. CartoonGAN has its heavily parameterized generator (56.46 GFLOPs on $256\times 256$ images) publicly available.Available at https://github.com/maciej3031/comixify. Since CartoonGAN has a feed-forward encoder-decoder structure, without using cycle consistent loss, CEC is not directly applicable to compress it. So we only compare GS with the other published state-of-the-art method GD on this task. Experiments are conducted on the CelebA dataset . Following , we use a student generator with $1/6$ channels of the teacher generator for GD, which achieves similar (but less) compression ratio compared with GS.

The visual results of cartoon style transfer, together with model statistics (FLOPs and model sizes), are shown in Fig. 5.Following , we use color matching as the post-processing on all compared methods, for better visual display quality. All FLOPs are calculated for input images with shape $256\times 256$. At large compression ratio, the style transfer results of GD have obvious visual artifacts (e.g., abnormal white spots). In contrast, GS-32 can remarkably compress the original generator by around $42\times$ (in terms of FLOPs) with minimal degradation in the visual quality. GS-8 can further improve the FLOPs compression ratio to $47\times$ with almost identical visual quality. These results again show the superiority of our student generator jointly learned by channel pruning, quantization and distillation, over the hand-crafted student generator used in GD. Part of the proposed GS framework is integrated into some style transfer products in Kwai Inc.’s Apps.

Conclusion

In this paper, we propose the first end-to-end optimization framework combining multiple compression techniques for GAN compression. Our method integrates model distillation, channel pruning and quantization, within one unified minimax optimization framework. Experimental results show that our method largely outperforms existing GAN compression options which utilize single compression techniques. Detailed ablation studies show that naively stacking different compression methods fails to achieve satisfying GAN compression results, sometimes even hurting the performance catastrophically, therefore testifying the necessity of our unified optimization framework.

References

Appendix 0.A Image Generation with SNGAN

We have demonstrated the effectiveness of GS in compressing image-to-image GANs (e.g., CycleGAN , StyleGAN ) in the main text. Here we show GS is also generally applicable to noise-to-image GANs (e.g., SNGAN ). SNGAN with the ResNet backbone is one of the most popular noise-to-image GANs, with state-of-the-art performance on a few datasets such as CIFAR10 [krizhevsky2009learning]. The generator in SNGAN has 7 convolution layers with 1.57 GFLOPs, with $32\times 32$ image outputs. We evaluate SNGAN generator compression on the CIFAR-10 dataset. Inception Score (IS) [salimans2016improved] is used to measure image generation and style transfer quality. We use latency (FLOPs) and model size to evaluate the network efficiency. Quantative and visualization results are shown in Table 3 and Figure 6 respectively. GS is able to compress SNGAN by up to $8\times$ (in terms of model size), with minimum drop in both visual quality and the quantitative IS value of generated images.