MobileStyleGAN.pytorch

MobileStylegan: una red neuronal convolucional ligera para la síntesis de imágenes de alta fidelidad

Implementación oficial de Pytorch

Los videos que lo acompañan se pueden encontrar en YouTube. Para obtener más detalles, consulte el documento.

Requisitos

• Python 3.8+
• 1–8 GPU NVIDIA de alta gama con al menos 12 GB de memoria. Hemos realizado todas las pruebas y desarrollo utilizando DL Workstation con 4x2080ti

Capacitación

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

Convertir el punto de control de Rosinity/stylegan2-pytorch

Nuestro marco admite el formato de puntos de control StyleGan2 de Rosinity/Stylegan2-Pytorch. Para convertir CKPT su propio punto de control de Stylegan2 a nuestro marco:

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 >

Verifique el punto de control convertido

Para verificar que su punto de control se convierta correctamente, simplemente ejecute la visualización de demostración:

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

Genere imágenes usando 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 > 

Evaluar la puntuación FID

Para evaluar la puntuación FID utilizamos una versión modificada de la biblioteca Pytorch-Fid:

python evaluate_fid.py < path_to_ref_dataset > < path_to_generated_imgs > 

Manifestación

Ejecute la visualización de demostración usando MobileStylegan:

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

Ejecute la comparación visual usando Stylegan2 vs. Mobilestylegan:

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

Convertir a ONNX

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

Convertir a coreMl

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

Implementación utilizando OpenVino

Proporcionamos la biblioteca externa Random_face como un ejemplo de implementación de nuestro modelo en los dispositivos Edge utilizando el marco OpenVino.

Modelos previos a la aparición

Nombre	DEFENSOR
mobilestylegan_ffhq.ckpt	7.75

(*) Nuestro marco admite modelos de detención de descarga automática, solo use --ckpt <pretrined_model_name> .

Licencia heredada

Código	Fuente	Licencia
Kernels CUDA personalizados	https://github.com/nvlabs/stylegan2	Licencia NVIDIA
Bloques de Stylegan2	https://github.com/rosinity/stylegan2-pytorch	MIT

Expresiones de gratitud

Queremos agradecer a las personas cuyas obras contribuyeron a nuestro proyecto ::

• Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila para investigaciones relacionadas con modelos generativos basados ​​en estilo.
• Kim Seonghyeon para la implementación de Stylegan2 en Pytorch.
• Cotter Fergal para la implementación de transformaciones de wavelet discretas y transformaciones wavelet discretas inversas en Pytorch.
• Cyril Diagne para la excelente demostración de cómo ejecutar Mobilestylegan directamente en el navegador web.

Citación

Si está utilizando los resultados y el código de este trabajo, cíquelo como:

 @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},
}