ID-disentanglement-Pytorch
Pytorch implementation of the paper Face Identity Disentanglement via Latent Space Mapping for both training and evaluation, with StyleGAN 2.
- Paper: https://arxiv.org/abs/2005.07728
- Official TensorFlow Implementation: https://github.com/YotamNitzan/ID-disentanglement
- Important: this implementation does not follow the exact same architecture of the original paper. See changes below
Changes from original paper
- instead of using a Discriminator loss for the mapper. We have used several other losses such as:
- LPIPS Loss (The Unreasonable Effectiveness of Deep Features as a Perceptual Metric, Zhang el al, 2018)
- MSE Loss
- Different ID Loss
- Different landmark detector
- The reason for those changes resides in the fact that the training procedure with Discriminator is often
hard and does not converge. We have found that replacing the Discriminator with LPIPS and MSE losses
we can achieve the same result. Nevertheless, our code supports training with a discriminator which can be
activated using the configuration. - The other changes are due to better Recognition models that have developed since the original paper was published
Setup
We used several pretrained models:
- StyleGan2 Generator for image size 256 - 550000.pt
- ID Encoder - model_ir_se50.pth
- Landmarks Detection - mobilefacenet_model_best.pth.tar
Weight files attached at this Drive folder.
You can also find at the above link our environment.yml file to create a relevant conda environment.
Training
The dataset is comprised of StyleGAN 2 generated images and W latent codes. see Utils/data_creator.py.
Examples of our generated dataset attached at this Drive folder.
To train the model run train_script.py, you can change parameters in Configs/ folder.
Inference
Try Inference.ipynb notebook to disentangle identity from attributes by yourself
Checkpoints
Our pretrained checkpoint attached at this Drive folder.
Results
