MobileStyleGAN.pytorch

MobileSylegan: un réseau neuronal convolutionnel léger pour la synthèse d'image haute fidélité

Implémentation officielle de Pytorch

Les vidéos d'accompagnement se trouvent sur YouTube. Pour plus de détails, veuillez vous référer au document.

Exigences

• Python 3.8+
• 1–8 GPU NVIDIA haut de gamme avec au moins 12 Go de mémoire. Nous avons effectué tous les tests et développement en utilisant DL Workstation avec 4x2080Ti

Entraînement

pip install -r requirements.txt
python train.py --cfg configs/mobile_stylegan_ffhq.json --gpus < n_gpus > 

Convertir le point de contrôle de la rosinalité / Stylegan2-Pytorch

Notre framework prend en charge le format de points de contrôle Stylegan2 de Rosinality / Stylegan2-Pytorch. Pour convertir CKPT votre propre point de contrôle de Stylegan2 en notre cadre:

python convert_rosinality_ckpt.py --ckpt < path_to_rosinality_stylegan2_ckpt > --ckpt-mnet < path_to_output_mapping_network_ckpt > --ckpt-snet < path_to_output_synthesis_network_ckpt > --cfg-path < path_to_output_config_json >

Vérifier le point de contrôle converti

Pour vérifier que votre point de contrôle est converti correctement, il suffit d'exécuter une visualisation de démonstration:

python demo.py --cfg < path_to_output_config_json > --ckpt " " --generator teacher

Générer des images à l'aide de mobilestylegan

python generate.py --cfg configs/mobile_stylegan_ffhq.json --device cuda --ckpt < path_to_ckpt > --output-path < path_to_store_imgs > --batch-size < batch_size > --n-batches < n_batches > 

Évaluer le score FID

Pour évaluer le score FID, nous utilisons une version modifiée de la bibliothèque Pytorch-Fid:

python evaluate_fid.py < path_to_ref_dataset > < path_to_generated_imgs > 

Démo

Exécutez une visualisation de démonstration à l'aide de MobileSylegan:

python demo.py --cfg configs/mobile_stylegan_ffhq.json --ckpt < path_to_ckpt >

Exécutez une comparaison visuelle en utilisant Stylegan2 vs MobileSylegan:

python compare.py --cfg configs/mobile_stylegan_ffhq.json --ckpt < path_to_ckpt > 

Convertir en onnx

python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt < path_to_ckpt > --export-model onnx --export-dir < output_dir > 

Convertir en coreml

python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt < path_to_ckpt > --export-model coreml --export-dir < output_dir > 

Déploiement à l'aide d'OpenVino

Nous fournissons une bibliothèque externe Random_face comme exemple de déploiement de notre modèle sur les périphériques Edge à l'aide du framework OpenVino.

Modèles pré-entraînés

Nom	Fid
mobilestylegan_ffhq.ckpt	7.75

(*) Notre framework prend en charge les modèles de téléchargement automatique de téléchargement, utilisez simplement --ckpt <pretrined_model_name> .

Licence héritée

Code	Source	Licence
Grains de cuda personnalisés	https://github.com/nvlabs/stylegan2	Licence Nvidia
Blocs Stylegan2	https://github.com/rosinality/stylegan2-pytorch	Mit

Remerciements

Nous tenons à remercier les personnes dont les œuvres ont contribué à notre projet ::

• Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila pour la recherche liée aux modèles génératifs basés sur le style.
• Kim Seonghyeon pour la mise en œuvre de Stylegan2 à Pytorch.
• Fergal Cotter pour la mise en œuvre de transformations en ondelettes discrètes et des transformations d'ondelettes discrètes inverses en pytorch.
• Cyril diagnostique pour l'excellente démo de la façon de gérer MobileSylegan directement dans le navigateur Web.

Citation

Si vous utilisez les résultats et le code de ce travail, veuillez le citer comme:

 @misc{belousov2021mobilestylegan,
      title={MobileStyleGAN: A Lightweight Convolutional Neural Network for High-Fidelity Image Synthesis},
      author={Sergei Belousov},
      year={2021},
      eprint={2104.04767},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

@article{BELOUSOV2021100115,
      title = {MobileStyleGAN.pytorch: PyTorch-based toolkit to compress StyleGAN2 model},
      journal = {Software Impacts},
      year = {2021},
      issn = {2665-9638},
      doi = {https://doi.org/10.1016/j.simpa.2021.100115},
      url = {https://www.sciencedirect.com/science/article/pii/S2665963821000452},
      author = {Sergei Belousov},
}