README.md

# StyleGAN2
## 论文
- https://arxiv.org/pdf/1912.04958
## 模型结构
针对StyleGAN中，大多数的生成的图像容易产生一个类似水滴状的伪影问题，StyleGAN2对AdaIN的归一化操作进行改进，作者把AdaIN层里面的归一化去掉，将噪声B和偏置项b移动出style模块之外，取得了更好的生成效果。

<div align=center>
    <img src="./assets/AdaIN.png"/>
</div>

## 算法原理
StyleGAN2为StyleGAN的改进版本，StyleGAN是在PGGAN的基础上进行改进的模型  
GAN（生成对抗网络）包含了一个生成模型G和一个判别模型D，模型通过生成器G对从正态分布P(z)随机采样的z生成伪数据x'，和从真实图像分布P(data)中采样的样本x作为判别器D的输入，判别器要让x的概率越大越好，让x'的概率越小越好，同时生成器希望生成的样本让判别器判别为真的概率越大越好。通过这种对抗的方式使模型生成越来越逼真的图片。

<div align=center>
    <img src="./assets/GAN2.png"/>
</div>  

PGGAN（渐进式生长生成对抗网络）通过先从低分辨率开始训练，然后再逐层提高分辨率进行训练的方式解决了传统GAN存在的模式崩溃（生成数据只是原始数据的子集（生成器偏向于生成判别器难以判别的样本））和难以训练高分辨率图片（生成器刚开始直接生成高分辨率图片很容易被判别器识别，在反向传播出现梯度大范围更新导致生成器崩溃）的问题，PGGAN先训练低分辨率，然后通过平滑接入的方式逐步提高分辨率：

<div align=center>
    <img src="./assets/PGGAN.png"/>
</div>  
 
PGGAN虽然能生成高清伪图，但是不能对图象的风格和细节进行修改，StyleGAN通过对特征进行解耦，让特征之间相互独立，互不影响，从而达到单独修改图象的某一部分的目的。具体来说StyleGAN是通过修改生成器，下图中左边为传统的生成器，右部分为StyleGAN的生成器，由两部分构成——Mapping network和Synthesis network ，其中Mapping network就是用来控制图像的风格信息，Synthesis network用来生成图像

<div align=center>
    <img src="./assets/styleGAN.png"/>
</div> 

StyleGAN2就是在StyleGAN的基础上改进了归一化操作，对损失函数和训练方法进行改进。


## 环境配置
### Docker（方法一）
从[光源](https://www.sourcefind.cn/#/service-list)中拉取docker镜像：
```
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.2-py3.10
```
创建容器并挂载目录进行开发：
```
docker run -it --name {name} --shm-size=1024G  --device=/dev/kfd --device=/dev/dri/ --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ulimit memlock=-1:-1 --ipc=host --network host --group-add video -v /opt/hyhal:/opt/hyhal:ro -v {}:{} {docker_image} /bin/bash
# 修改1 {name} 需要改为自定义名称，建议命名{框架_dtk版本_使用者姓名},如果有特殊用途可在命名框架前添加命名
# 修改2 {docker_image} 需要需要创建容器的对应镜像名称，如: image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.2-py3.10
# 修改3 -v 挂载路径到容器指定路径
pip install -r requirements.txt
```
### Dockerfile（方法二）
```
cd docker
docker build --no-cache -t gan2_pytorch:1.0 .
docker run -it --name {name} --shm-size=1024G  --device=/dev/kfd --device=/dev/dri/ --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ulimit memlock=-1:-1 --ipc=host --network host --group-add video -v /opt/hyhal:/opt/hyhal:ro -v {}:{} {docker_image} /bin/bash 
pip install -r requirements.txt
```
### Anaconda（方法三）
线上节点推荐使用conda进行环境配置。
创建python=3.10的conda环境并激活
```
conda create -n styleGan2 python=3.10
conda activate styleGan2
```

关于本项目DCU显卡所需的特殊深度学习库可从[光合](https://developer.hpccube.com/tool/)开发者社区下载安装。
```
DTK驱动：dtk24.04.2
python：python3.10
pytorch:2.1.0
torchvision:0.16.0
```
安装其他依赖包
```
pip install -r requirements.txt
```
## 数据集

- unconditional models训练集[FFHQ](https://drive.google.com/drive/folders/1u2xu7bSrWxrbUxk-dT-UvEJq8IjdmNTP)，

- 用于快速复现的小数据集SCNet快速下载链接 [images](http://113.200.138.88:18080/aidatasets/project-dependency/stylegan2_pytorch)

下载后解压到当前目录  
数据目录结构如下：
```
./data/ffhq/images
    └── 00000.png
    └── 00001.png
    └── ...
```


## 训练

先将数据处理为lmdb格式：  
```
python prepare_data.py --out LMDB ./data/ffhq/images
```

### 单机单卡
```
HIP_VISIBLE_DEVICES=2 python train.py LMDB --batch 4
```

### 单机多卡
```
HIP_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 --master_port=29500 train.py --batch 4 LMDB
```

## 推理
下载权重文件[stylegan2-ffhq-config-f.pt](https://pan.baidu.com/s/1IOD4DjjjrMZBF_TfL7pGJQ?pwd=1234)  
SCNet快速下载链接：[stylegan2-ffhq-config-f.pt](http://113.200.138.88:18080/aimodels/findsource-dependency/stylegan2_pytorch)  

模型推理：
```
python generate.py --sample 6 --pics 6 --ckpt stylegan2-ffhq-config-f.pt
```
推理结果：  
<div align=center>
    <img src="./assets/000000.png"/>
</div> 


## 应用场景   
### 算法类别  
人脸生成  

### 热点应用行业
安防,交通,教育


## 源码仓库及问题反馈
[https://developer.hpccube.com/codes/modelzoo/stylegan2_pytorch](https://developer.hpccube.com/codes/modelzoo/stylegan2_pytorch)
## 参考资料
[https://github.com/rosinality/stylegan2-pytorch](https://github.com/rosinality/stylegan2-pytorch)