Skip to content

Implementation of VAE and Style-GAN Architecture Achieving State of the Art Reconstruction

License

Notifications You must be signed in to change notification settings

orgoro/style-vae

Repository files navigation

A Style Based Variational Autoencoder

Build Status

Abstract

VAE are among the state of the art generative models, but have recently lost their shine to GANs. The most prominent work recently in which is the Style-GAN by Karras et al. VAE has the ability to encode as well as decode - this advantage over the style-gan is useful in many downstream tasks. In this work we combine the style based architecture and VAE and achieve state of the art reconstruction and generation. We follow the work of Hou et al. DFC-VAE to use perceptual loss and we compare our results to this work.

Dataset

  • Celeb-A - the results below model (64x64)
  • FFHQ - for other models (128x128, 256x256)

Architecture

arch

The loss is comprised out of two components:

  • Reconstruction Loss - based on perceptual loss (pre-trained VGG16 features)
  • Latent Loss - kl-divergence loss

Reconstruction Results

recon

Random Sample Results

random

Interpolation

random

Comparison to State of the Art

compare

Code Structure

Training Curves

curve

Test

Test results can be seen in the visuals part

Additional Experiments

We used the provided model to train on the FFHQ dataset to produce a 256x256 results: 64-256

Future Work

  • Adversarial Training- Adding an adversarial term to improve generation of hair and fine details (hair) and background especially in the high resulotion models
  • Weighted KL-Divergence- The hierarchical structure of the code injection would allow weighted KL-divergence loss term where we allow the VAE encoder to different divergence fine/coarse features. This flexibility makes sense perhaps because fine details like hairs are more normally distributed, while coarse feature behave differently (not many dark skin red heads i.e.)

Trained Models

Available soon...

About

Implementation of VAE and Style-GAN Architecture Achieving State of the Art Reconstruction

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages