Towards Device Efficient Conditional Image Generation

TL;DR

This paper introduces a device-agnostic, two-stage autoencoder optimization method that reduces computational requirements for conditional image generation without losing image quality, enabling real-time deployment on CPU devices.

Contribution

The authors propose a novel channel condensation and pruning strategy combined with student-teacher fine-tuning to optimize autoencoders for various devices, maintaining quality while reducing compute.

Findings

Achieved real-time autoencoders on CPU-only devices.

Demonstrated performance improvements across multiple image generation tasks.

Validated effectiveness through ablation studies.

Abstract

We present a novel algorithm to reduce tensor compute required by a conditional image generation autoencoder without sacrificing quality of photo-realistic image generation. Our method is device agnostic, and can optimize an autoencoder for a given CPU-only, GPU compute device(s) in about normal time it takes to train an autoencoder on a generic workstation. We achieve this via a two-stage novel strategy where, first, we condense the channel weights, such that, as few as possible channels are used. Then, we prune the nearly zeroed out weight activations, and fine-tune the autoencoder. To maintain image quality, fine-tuning is done via student-teacher training, where we reuse the condensed autoencoder as the teacher. We show performance gains for various conditional image generation tasks: segmentation mask to face images, face images to cartoonization, and finally CycleGAN-based model over multiple compute devices. We perform various ablation studies to justify the claims and design choices, and achieve real-time versions of various autoencoders on CPU-only devices while maintaining image quality, thus enabling at-scale deployment of such autoencoders.