MobileStyleGAN.pytorch

Mobilestylegan：高保真圖像綜合的輕量級卷積神經網絡

Pytorch官方實施

隨附的視頻可以在YouTube上找到。有關更多詳細信息，請參閱論文。

要求

• Python 3.8+
• 1–8高端NVIDIA GPU，至少有12 GB的內存。我們已經使用DL工作站進行了所有測試和開發

訓練

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

轉換來自Rosinality/stylegan2-Pytorch的檢查點

我們的框架支持Rosinality/stylegan2-Pytorch的stylegan2檢查點格式。要將您自己的stylegan2檢查站轉換為我們的框架：

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 >

檢查轉換的檢查點

要檢查檢查點是否正確轉換，只需運行演示可視化：

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

使用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 > 

評估FID分數

為了評估FID分數，我們使用Pytorch-FID庫的修改版本：

python evaluate_fid.py < path_to_ref_dataset > < path_to_generated_imgs > 

演示

使用Mobilestylegan運行演示可視化：

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

使用stylegan2 vs. mobilestylegan進行視覺比較：

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

轉換為ONNX

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

轉換為Coreml

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

使用OpenVino部署

我們提供外部庫Random_face作為使用OpenVino框架在邊緣設備上部署模型的示例。

預驗證的模型

姓名	fid
MobileStyLegan_FFHQ.CKPT	7.75

（*）我們的框架支持自動下載預處理的型號，只需使用--ckpt <pretrined_model_name>即可。

遺產許可證

代碼	來源	執照
自定義CUDA內核	https://github.com/nvlabs/stylegan2	NVIDIA許可證
stylegan2塊	https://github.com/rosinality/stylegan2-pytorch	麻省理工學院

致謝

我們要感謝那些為我們的項目做出貢獻的人們::

• Tero Karras，Samuli Laine，Miika Aittala，Janne Hellsten，Jaakko Lehtinen，Timo Aila，用於與基於樣式的生成模型有關的研究。
• Kim Seonghyeon用於在Pytorch中實施stylegan2。
• 用於實施離散小波變換的Fergal Cotter和Pytorch中的逆離散小波變換。
• Cyril Diagne提供瞭如何直接運行Mobilestylegan進入Web瀏覽器的出色演示。

引用

如果您正在使用此工作的結果和代碼，請將其引用為：

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