MobileStyleGAN：用于高保真图像合成的轻量级卷积神经网络

官方 PyTorch 实现

配套视频可在 YouTube 上找到。 更多详情请参阅 论文。

环境要求

• Python 3.8+
• 1–8 张高端 NVIDIA GPU，每张显存至少 12 GB。我们所有的测试和开发均在配备 4 块 2080 Ti 显卡的深度学习工作站上完成。

训练

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 <rosinality_stylegan2_ckpt路径> --ckpt-mnet <输出映射网络检查点路径> --ckpt-snet <输出合成网络检查点路径> --cfg-path <输出配置文件JSON路径>

检查转换后的检查点

要确认您的检查点已正确转换，只需运行演示可视化：

python demo.py --cfg <输出配置文件JSON路径> --ckpt "" --generator teacher

使用 MobileStyleGAN 生成图像

python generate.py --cfg configs/mobile_stylegan_ffhq.json --device cuda --ckpt <检查点路径> --output-path <保存图像路径> --batch-size <批量大小> --n-batches <批次数量>

评估 FID 分数

我们使用 pytorch-fid 库的修改版本来评估 FID 分数：

python evaluate_fid.py <参考数据集路径> <生成图像路径>

演示

使用 MobileStyleGAN 运行演示可视化：

python demo.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <检查点路径>

使用 StyleGAN2 和 MobileStyleGAN 进行视觉对比：

python compare.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <检查点路径>

转换为 ONNX

python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <检查点路径> --export-model onnx --export-dir <输出目录>

转换为 CoreML

python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <检查点路径> --export-model coreml --export-dir <输出目录>

使用 OpenVINO 部署

我们提供外部库 random_face，作为使用 OpenVINO 框架在边缘设备上部署我们模型的示例。

预训练模型

名称	FID
mobilestylegan_ffhq.ckpt	7.75

(*) 我们的框架支持自动下载预训练模型，只需使用 --ckpt <预训练模型名称> 即可。

旧版许可证

代码	来源	许可证
自定义 CUDA 内核	https://github.com/NVlabs/stylegan2	Nvidia 许可证
StyleGAN2 块	https://github.com/rosinality/stylegan2-pytorch	MIT

致谢

我们感谢为本项目做出贡献的人们：

• Tero Karras、Samuli Laine、Miika Aittala、Janne Hellsten、Jaakko Lehtinen、Timo Aila，感谢他们在基于风格的生成模型方面的研究。
• Kim Seonghyeon，感谢他在 PyTorch 中实现 StyleGAN2。
• Fergal Cotter，感谢他在 PyTorch 中实现离散小波变换及逆离散小波变换。
• Cyril Diagne，感谢他出色的 MobileStyleGAN 直接在网页浏览器中运行的演示。

引用

如果您使用了本项目的成果和代码，请按以下方式引用：

@misc{belousov2021mobilestylegan,
      title={MobileStyleGAN：用于高保真图像合成的轻量级卷积神经网络},
      author={Sergei Belousov},
      year={2021},
      eprint={2104.04767},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

@article{BELOUSOV2021100115,
      title = {MobileStyleGAN.pytorch：基于 PyTorch 的 StyleGAN2 模型压缩工具},
      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},
}

MobileStyleGAN.pytorch 快速上手指南

MobileStyleGAN 是一个轻量级的卷积神经网络，旨在实现高保真度的图像合成，特别适合在移动端或边缘设备部署。本指南将帮助你快速搭建环境并生成图像。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.8 或更高版本
• GPU: 1–8 张高端 NVIDIA 显卡（显存至少 12 GB）。官方测试环境为 4x RTX 2080 Ti。
• 依赖库: 需要安装 PyTorch 及相关深度学习库。

提示：国内用户建议使用清华源或阿里源加速 Python 包的安装。

安装步骤

1. 克隆项目代码

   git clone https://github.com/bes-dev/MobileStyleGAN.pytorch.git
   cd MobileStyleGAN.pytorch
   

2. 安装依赖 使用 pip 安装所需依赖包。国内用户可添加 -i 参数指定镜像源以提升下载速度：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

1. 下载预训练模型

本项目支持自动下载预训练模型。你可以直接使用模型名称作为 --ckpt 参数，或者手动下载 mobilestylegan_ffhq.ckpt 并放置于项目目录。

2. 生成图像

使用以下命令即可基于预训练模型生成高质量人脸图像：

python generate.py --cfg configs/mobile_stylegan_ffhq.json --device cuda --ckpt mobilestylegan_ffhq.ckpt --output-path ./generated_images --batch-size 4 --n-batches 10

• --cfg: 配置文件路径。
• --device: 运行设备（cuda 或 cpu）。
• --ckpt: 检查点路径（可直接填模型名自动下载，或填本地路径）。
• --output-path: 生成图像的保存目录。
• --batch-size 和 --n-batches: 控制生成的图片数量。

3. 交互式演示 (可选)

如果你想直观地查看生成效果或对比 StyleGAN2 与 MobileStyleGAN 的差异，可以运行演示脚本：

# 运行 MobileStyleGAN 可视化演示
python demo.py --cfg configs/mobile_stylegan_ffhq.json --ckpt mobilestylegan_ffhq.ckpt

# 对比 StyleGAN2 与 MobileStyleGAN
python compare.py --cfg configs/mobile_stylegan_ffhq.json --ckpt mobilestylegan_ffhq.ckpt

4. 模型导出 (可选)

若需部署到移动端或边缘设备，可将模型导出为 ONNX 或 CoreML 格式：

# 导出为 ONNX
python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt mobilestylegan_ffhq.ckpt --export-model onnx --export-dir ./onnx_model

# 导出为 CoreML
python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt mobilestylegan_ffhq.ckpt --export-model coreml --export-dir ./coreml_model