MobileStyleGAN.pytorch

Mobilestylegan: uma rede neural convolucional leve para a síntese de imagem de alta fidelidade

Implementação oficial de Pytorch

Os vídeos que acompanham podem ser encontrados no YouTube. Para mais detalhes, consulte o artigo.

Requisitos

• Python 3.8+
• 1–8 GPUs NVIDIA de ponta com pelo menos 12 GB de memória. Fizemos todos os testes e desenvolvimento usando DL Workstation com 4x2080ti

Treinamento

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

Converta o ponto de verificação da rosinalidade/stylegan2-pytorch

Nossa estrutura suporta o formato Stylegan2 Points do rosinalidade/stylegan2-pytorch. Para converter o CKPT seu próprio ponto de verificação do Stylegan2 em nossa estrutura:

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 o ponto de verificação convertido

Para verificar se seu ponto de verificação é convertido corretamente, basta executar a visualização de demonstração:

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

Gerar imagens 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 > 

Avalie a pontuação do FID

Para avaliar a pontuação do FID, usamos uma versão modificada da Biblioteca Pytorch-FID:

python evaluate_fid.py < path_to_ref_dataset > < path_to_generated_imgs > 

Demonstração

Execute a visualização de demonstração usando o Mobilestylegan:

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

Execute a comparação visual usando Stylegan2 vs. Mobilestylegan:

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

Converter para onnx

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

Converter em coreml

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

Implantação usando o OpenVino

Fornecemos biblioteca externa random_face como um exemplo de implantação de nosso modelo nos dispositivos Edge usando a estrutura OpenVino.

Modelos pré -teremam

Nome	Fid
mobilestylegan_ffhq.ckpt	7.75

(*) Nossa estrutura suporta modelos de download automático pré -terenciado, basta usar --ckpt <pretrined_model_name> .

Licença Legada

Código	Fonte	Licença
Kernels CUDA personalizados	https://github.com/nvlabs/stylegan2	Nvidia Licença
Blocos Stylegan2	https://github.com/rosinality/stylegan2-pytorch	Mit

Agradecimentos

Queremos agradecer às pessoas cujas obras contribuíram para o nosso projeto ::

• Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila para pesquisas relacionadas a modelos generativos baseados em estilo.
• Kim Seonghyeon para implementação de Stylegan2 em Pytorch.
• Cotter Fergal para implementação de transformações discretas de wavelet e transformações inversas de wavelets discretos em Pytorch.
• Cyril Diagna para a excelente demonstração de como administrar o Mobilestylegan diretamente no navegador da web.

Citação

Se você estiver usando os resultados e o código deste trabalho, cite -o 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},
}