Real-Time Portrait Stylization on the Edge

1 Introduction

Thanks to hardware advancement, varieties of AI applications have been enabled on portable smart devices, such as foreground segmentation, face recognition, etc. In this work, we investigate portrait stylization which is a popular feature in social media Apps, transferring self portraits into a desired style, such as cartoon [Andersson and Arvidsson, 2020], anime [Kim et al., 2020; Li et al., 2021], etc.

Portrait stylization can be categorized as classic imageto-image translation, which is often achieved by conditional Generative Adversarial Networks (GANs) [Isola et al., 2017; Zhu et al., 2017]. In GAN training, a generator learns to generate a fake instance and fool the discriminator, while the discriminator takes true and generated images as input and learns to distinguish them. Consequently, in the portrait stylization task, the generator is utilized to map a portrait photo to the desired style domain during inference. Compared to naive paired supervised training, GANs exhibits stunning generative quality, e.g., sharp and realistic details, diverse features.

However, it still remains challenging to enable real-time face stylization to process videos on a mobile device, and the reason comes two-fold. Firstly, typically following an encoder-decoder design, image translation models suffer from high computation complexity, especially on highresolution images. Secondly, GAN training is difficult and unstable, suffering from loss divergence and mode collapse. As a result, existing compression techniques are difficult to integrate into GAN training and preserve generative quality.

In this work, we propose a compiler-aware differentiable architecture search framework. We measure the latency of the building blocks with sufficient configurations (channels, feature sizes), and train a neural network to predict the latency. We show that a simple MLP speed model can make accurate predictions. In order to search compact architectures, we integrate learnable parameters in the GAN generator and regularize them by speed constraints to reduce the model width and depth. With the speed model that maps architecture parameters to latency, the speed penalty is differentiable so that we can perform search with end-to-end training. Plus, different from prior work, we do not select preserved/eliminated model weights by magnitude. Instead, we apply straight through estimator on architecture parameters to sparsify them into {0, 1}. The benefit is two-fold. The remained weights are represented by 1s so that we can easily predict the latency of a certain state. In addition, the gradients of pruned weights are completely zeroed out so that their functionalities are preserved. As a result, pruned weights are always ready to be reverted back to contribute to accuracy during exploration. This is especially important in GAN search because of the unstable training process.

∗These authors contributed equally to this work.

Overall, our contributions include:

• We develop a latency-driven differentiable architecture search for GANs. Our sophisticated pipeline addresses the difficulty in GAN searching, achieves unprecedented compression rate while preserving generative quality.

• To the best of our knowledge, we are the first to achieve real-time portrait stylization on mobile phones. Mobile demos are attached in the link 1 .

2 Background

Image-to-Image Style Transfer. [Isola et al., 2017; Liu et al., 2017; Huang et al., 2018; Lee et al., 2018; Liu et al., 2019; Park et al., 2020; Kim et al., 2020; Li et al., 2021; Chong and Forsyth, 2021; Tang et al., 2019] aim to translate a image from the source domain to match certain styles in the target domain, while preserving semantics of the origin image. Early works train the generative model with paired data [Isola et al., 2017], however they cannot be applied to more common unpaired datasets. Later work [Zhu et al., 2017] proposed a cycle consistency loss to train on unpaired data domains, inspiring lots of successor research on image stylization [Park et al., 2020; Kim et al., 2020; Li et al., 2021; Chong and Forsyth, 2021; Tang et al., 2021].

As for the domain of cartoon or anime, UGATIT [Kim et al., 2020] released the selfie2anime benchmark and develop an adaptive mixture of instance and layer normalization. AniGAN [Li et al., 2021] released a larger scale face2anime dataset and studied style-guided anime translation. GNR [Chong and Forsyth, 2021] refine the content and style to produce controllable and diverse synthesis.

Compressing GANs. Because of wide applications, the compression of GAN has drawn research attention [Fu et al., 2020; Wang et al., 2020; Li et al., 2020; Chen et al., 2021; Jin et al., 2021]. Recent work [Li et al., 2020] proposed to incorporate neural architecture search (NAS) and feature level knowledge distillation (KD). [Jin et al., 2021] integrated an Inception-like residual block and performed self-distillation.

3 Compiler-Aware Architecture Search

3.1 GAN Basics

We follow CycleGAN [Zhu et al., 2017] and UGATIT [Kim et al., 2020] paradigm to develop our architecture search for stylization. According to CycleGAN, we learn mapping functions between two unpaired domain X={xi}Ni=1 and Y={yj}Mj=1 . There are two generators in inverse direction G:X-→Y and F:Y-→X, as well as two discriminators DX and DY . Note that we refer G as the generator mapping portraits to stylized images, which is the only required model during inference.

Adversarial Loss. We match the distribution of translated instances with the target domain as follows.

CycleConsistencyLoss. Tominimize reconstruction error,

Note that we also incorporate the identity loss Lid and CAM loss LCAM as proposed in UGATIT, for simplicity we skip detailed formulations here and please refer to [Kim et al., 2020].

1 The link for the demo video: https://youtu.be/2SFAVwaymvQ

Our overall GAN objective is:

where λ1=1, λ2=10, λ3=10, λ4=1000 are the hyperparameters to control each loss.

3.2 LayerwiseWidth Search

Width search is performed for each CONV layer. We choose the supernet from [Zhu et al., 2017; Kim et al., 2020], which is a commonly employed generator. In order to create a learnable binary mask, we insert a depth-wise 1×1 CONV layer following each CONV layer to be pruned, as shown below,

where ⊙ denotes the convolution operation. wn l ∈Ro×i×k×k is the weight parameters in the lth CONV layer of the nth block, with o output channels, i input channels, and kernels of size k×k. anl ∈RB×o×s×s′ represents the output features of lth layer (with the trainable mask), with o channels and s×s′ feature size. B denotes the batch size. vn l ∈Ro×1×1×1 is the corresponding weights of the depth-wise CONV layer (i.e., the mask layer).

Larger elements of mnl mean that the corresponding channels should be preserved while smaller elements indicate that the corresponding channels should be pruned. Formally, we use a threshold t to convert mnl into a binary mask as below,

where bnl ∈{0, 1}o×1×1×1 is the binarized vn l . Typically, we initialize vn l with 1, and the adjustable t is set to 0.5. In order to make the mask differentiable to enable backpropagation, we utilize Straight Through Estimator (STE) [Bengio et al., 2013; Chang et al., 2020] as shown below,

Our trainable binary mask has the following advantages: (i) The mask can be trained along with the network parameters via gradient descent, thus saving search cost compared to NAS methods [Zoph and Le, 2017; Zhong et al., 2018]. (ii) Different from previous methods [Han et al., 2015; Yu et al., 2017; He et al., 2017; Guan et al., 2020], which determine the pruning according to the parameter magnitudes, we decouple the parameter magnitudes of CONV or BN layer from pruning, and utilize independent mask layers, thus the remained parameters are not harmed. (iii) The discrete values can directly provide the width information for each CONV layer, which is compatible with speed prediction.

3.3 Length Search by Block

Note that although per-layer width search may also converge to zero width, which eliminates the entire block, we find that there are usually a few channels left in each block preventing us to remove the entire block. Plus, the pruning indicator for each layer cannot represent the latency reduction of entirely pruned blocks. Thus it is necessary to perform length search separately from width search.

We construct two paths in each customized residual block, one is the masked convolution block and the other is skip connection. In the aggregation layer of the dual paths, we integrate binarized variables β1 and β2, then the forward computation can be represented as follows,

The aggregation layer contains two trainable parameters α1 and α2, shares similar STE recipe with width parameter. In the forward pass, it selects the skip path or the masked convolution path based on the relative relationship of α1 and α2,

3.4 Speed Prediction with SpeedModel

We take inference speed on mobile GPUs to constrain the optimization. A DNN-based speed model is adopted to predict the inference speed of the block based on its architecture configurations. Then the final predicted latency is accumulated by the aggregated blocks so that we can compute the latency loss Llatency and integrate it into the searching pipeline.

The trained speed model is accurate in predicting the speed of different layer widths in the block (with 5% error at most).

4 Experiments and Demonstration

We conduct experiments on Face2anime dataset published by [Li et al., 2021]. Face2anime consists of 17,796 images. We set input size to 256×256. Further, we also demonstrate realistic quality on cartoon stylization released by MiniVision, along with Asian women face training data generated by [Karras et al., 2020].

Experiment Setups. As for the supernet, we search from [Kim et al., 2020]. Learning rate is set to 2×10-4 for both generator and discriminator, with Adam optimizer. Leaning rate is fixed for the first 30k iterations and then linearly decayed to zero in another 30k iterations.

Performance Evaluation. We quantitatively compare our searched stylization model with representative works in the quality of created images (FID) and computation costs (MACs). As shown in Table 1, our generative model achieves 10× computation reduction and preserves generative quality. Thanks to the significant computation reduction, we achieve high quality real-time stylization on mobile.

Results Visualization and On-mobile Demonstration. Figure 2 shows the visualization of the created images on face2anime test dataset and the comparison of our 5.56 GMACs model with the baseline UGATIT model with 57.1 GMACs. Figure 3 shows the created images using photo2cartoon test dataset. With simpler styles, our search method further reduce the model size to 1.34 GMACs. We also demonstrate our method on the mobile device, as shown in Figure 1. The full demo video is available in the link 1 .

Acknowledgments

This research is partially supported by National Science Foundation CNS-1909172, CCF-1919117, and CCF2047516 (CAREER).