Awesome

PCA-based knowledge distillation towards lightweight and content-style balanced photorealistic style transfer models

Code for our CVPR 2022 paper PCA-based knowledge distillation towards lightweight and content-style balanced photorealistic style transfer models

Advantages of our distilled models

• They achieve a better balance between content preservation and style transferral. Specifically, they transfer stronger style effects than WCT2 and preserves better content than PhotoWCT and PhotoWCT2.
• They are lightweight. The models distilled from VGG and MobileNet use only 283K and 73K parameters, respectively. 73K is only 0.72% of the number of parameters in WCT2, 10.12M.
• They run fast. Our VGG-distilled model runs around 10x, 10x, and 5x as fast as WCT2, PhotoWCT, and PhotoWCT2. Our MobileNet-distilled model is about 2x as fast as Our VGG-distilled model.
• They can support the 8K (7680 x 4320) resolution on a normal commercial GPU.
• To our best knowledge, this is the first distillation for photorealistic style transfer. The closest work to ours is CKD for artistic style transfer.

Models and files

We apply our PCA distillation to two backbones: VGG and MobileNet. The corresponding files can be found in the folders VGG backbone and MobileNet backbone. For each backbone, we provide our trained parameters in the folder ckpts, the demo of how to perform style transfer in the file demo.ipynb, the training code in two files train_eigenbases.py and train_distilled_model.py, and the distilled model structure in utils/lightweight_model.py.

Note that train_eigenbases.py has to be run first to derive the global eignebases and train_distilled_model.py uses the eigenbases to distill models.

Requirements

• tensorflow v2.4.1 or above (we developed the models with tf-v2.4.1)

Citation

If you find this repo useful, please cite our paper PCA-based knowledge distillation towards lightweight and content-style balanced photorealistic style transfer models published in CVPR 2022.

This codebase is largely extended from our previous work PhotoWCT2: Compact Autoencoder for Photorealistic Style Transfer Resulting from Blockwise Training and Skip Connections of High-Frequency Residuals published in WACV2022. Please take a look at it if interested.