README.md

# stylegan3

## 论文

**Alias-Free Generative Adversarial Networks**  

* https://nvlabs-fi-cdn.nvidia.com/stylegan3/stylegan3-paper.pdf


## 模型结构

下图为生成器的结构，`Mapping network`（将Latent code z转换为w），`Fourier feat.`（傅里叶特征），`ToRGB`（将图像转换为RGB-3通道格式）。

![Alt text](readme_imgs/image-1.png)


## 算法原理

用途：该算法可以生成高质量图像。

原理：

![Alt text](readme_imgs/image-2.png)

以信号处理方法分析现有问题，并通过修改网络结构，在保证FID的前提下增加网络的等变性，具体操作如下，

1.使用傅里叶特征+删除噪声输入+删除跳跃链接+减少网络层数+禁用混合正则化和路径长度正则化

保持FID与stylegan2相似，同时略微提升平移等变性。

2.增加边距+sinc滤波器

增加平移等变性，但FID变差。

3.在上/下采样之间添加非线性函数

增加平移等变性。

4.降低截止频率

增加平移等变性，FID低于stylegan2。

5.Transformed Fourier features

改善FID。

6.Flexible layer specifications

提高等变性的质量。

7.旋转等变性

提高旋转等变性，降低可训练参数量（通过将3x3卷积替换为1x1卷积）。


## 环境配置

### Docker（方法一）

    docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.10
    docker run --shm-size 10g --network=host --name=stylegan3 --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal -it <your IMAGE ID> bash
    pip install -r requirements.txt

### Docker（方法二）

    # 需要在对应的目录下
    docker build -t <IMAGE_NAME>:<TAG> .
    # <your IMAGE ID>用以上拉取的docker的镜像ID替换
    docker run -it --shm-size 10g --network=host -v /opt/hyhal:/opt/hyhal --name=stylegan3 --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined <your IMAGE ID> bash
    pip install -r requirements.txt

### Anaconda (方法三)
1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装：
https://developer.sourcefind.cn/tool/

    DTK驱动：dtk24.04.1
    python：python3.10
    torch:2.1
    torchvision:0.16.1
    torchaudio:2.1.2
    deepspeed:0.12.3
    apex:1.1.0

Tips：以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应

2、其它非特殊库参照requirements.txt安装

    pip install -r requirements.txt


## 数据集

- 数据集快速下载中心:
- 官方下载地址
  - https://github.com/NVlabs/ffhq-dataset
  - https://github.com/NVlabs/metfaces-dataset
  - https://github.com/clovaai/stargan-v2/blob/master/README.md#animal-faces-hq-dataset-afhq

注意：训练的时候需要将一整个文件夹的数据转换成tfrecords的格式，可以通过以下命令生成对应的zip包 

    #例:处理AFHQv2数据集
    python dataset_tool.py --source=downloads/afhqv2 --dest=datasets/afhqv2-512x512.zip


## 推理

### 模型下载
模型权重快速下载中心
- 官方下载地址：
  - https://catalog.ngc.nvidia.com/orgs/nvidia/teams/research/models/stylegan3

```
pretrained_models/
    ├── xxx.pkl
    └── stylegan3-r-ffhq-1024x1024.pkl
```
注意：上述地址中包含多个模型，可以有选择地下载。

### 命令

    # 生成图像
    python gen_images.py --outdir=out --trunc=1 --seeds=2 --network=pretrained_models/stylegan3-r-ffhq-1024x1024.pkl

    # 生成视频
    python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 --network=pretrained_models/stylegan3-r-ffhq-1024x1024.pkl

注意：`network`既可以是本地模型也可以是url（如，`https://catalog.ngc.nvidia.com/orgs/nvidia/teams/research/models/stylegan3/files/xxx.pkl`）。下载的网络pickle缓存在`$HOME/.cache/dnnlib`，可以通过设置`DNNLIB_CACHE_DIR`环境变量来覆盖。默认的PyTorch扩展构建目录是`$HOME/.cache/torch_extensions`，可以通过设置`TORCH_EXTENSIONS_DIR`来覆盖。

## 训练

    # 训练新的网络
    # 例:训练 StyleGAN3-T 数据集AFHQv2, 4张DCU.
    python train.py --outdir=training-runs --cfg=stylegan3-t --data=datasets/afhqv2-512x512.zip \
    --gpus=4 --batch=8 --gamma=8.2 --mirror=1 --mbstd-group=2

    # 微调
    python train.py --outdir=training-runs --cfg=stylegan3-r --data=datasets/afhqv2-512x512.zip --gpus=1 --batch=1 --gamma=6.6 --mirror=1 --mbstd-group=1 --kimg=5000 --snap=5 --resume=stylegan3-r-afhqv2-512x512.pkl

注意：必须显式指定最重要的选项(`--gpu`、`--batch`和`--gamma`)，并且应该谨慎设置，请参阅`python train.py --help`，以获得完整的选项列表和一般指南和建议的训练配置，以及不同场景下预期的训练速度和内存使用情况。

每次训练运行的结果都保存到新创建的目录中，例如`training-runs/00000-stylegan3-t-afhqv2-512x512-gpus8-batch32-gamma8.2`。训练循环定期(由`--snap`控制)导出网络pickles (`network-snapshot-<king>.pkl`)和随机图像网格(`fakes<king>.png`)。对于每个导出的pickle，评估FID(由`--metrics`控制)并在`metric-fid50k_full.jsonl`中记录结果，同时还在`training_stats.jsonl`中记录各种统计数据。


## result

![image-20230710170324231](./pngs/show.gif)

### 精度

![image-20230710170324239](./pngs/loss.png)

## 应用场景

### 算法类别

`AIGC`

### 热点应用行业

`媒体,科研,教育`

## 源码仓库及问题反馈

* https://developer.sourcefind.cn/codes/modelzoo/stylegan3_pytorch

## 参考资料

* https://github.com/NVlabs/stylegan3