## [Alias-Free Generative Adversarial Networks (StyleGAN3)](https://nvlabs.github.io/stylegan3/)



### 模型介绍

StyleGAN3尝试解决生成模型（GAN）的一个普遍问题：生成的过程并不是一个自然的层次化生成，粗糙特征（GAN的浅层网络的输出特征）主要控制了精细特征（GAN的深层网络的输出特征）的存在与否，没有严格控制他们出现的精确位置。产生这个现象的根本原因是目前的生成器网络架构是卷积+非线性+上采样等结构，而这样的结构没有做到很好的 Equivariance（等变性），本文作者们就是要改进现有的生成器网络结构（主要针对StyleGAN2），使得其具有高质量的Equivariance。

StyleGAN3从信号处理的理论角度重新审视了生成模型框架，并提供了一种全新的生成模型设计方案来赋予网络Equivariance。



### 模型结构

主要针对生成器（StyleGAN2）结构进行改进，整个生成器网络的改进操作汇总为如下表格：

![image-20230629154853155](./pngs/stru.png)



### 数据集

下载[**FFHQ**](https://github.com/NVlabs/ffhq-dataset)、**[MetFaces](https://github.com/NVlabs/metfaces-dataset)**、[**AFHQv2**](https://github.com/clovaai/stargan-v2/blob/master/README.md#animal-faces-hq-dataset-afhq)或者准备自己的数据集

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

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



### 训练及推理

#### 环境配置

- 支持Linux和Windows，但出于性能和兼容性的原因，我们推荐Linux

- 1–4张 DCU卡

- 64-bit Python 3.8、dtk23.04，并下载对应版本的torch

  pytorch whl包下载目录：[https://cancon.hpccube.com:65024/4/main/pytorch/dtk23.04](https://cancon.hpccube.com:65024/4/main/pytorch/dtk23.04)

  ```shell
  pip install torch* (下载的torch的whl包)
  ```

- 其他包依赖：

  ```shell
  pip install -r requirements.txt
  ```

- GCC 7或更高版本(Linux)或Visual Studio (Windows)编译器、



#### 推理

预训练的网络存储为`*pkl`,可以使用本地或url的`pkl`文件

预训练模型地址：[StyleGAN3 pretrained models](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/research/models/stylegan3)

> 也可通过以下方式访问个别网络
>
> https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions`/1/files/<MODEL>`, `<MODEL>` 可以选:
>
> `stylegan3-t-ffhq-1024x1024.pkl`, `stylegan3-t-ffhqu-1024x1024.pkl`,`stylegan3-t-ffhqu-256x256.pkl`
>
> `stylegan3-r-ffhq-1024x1024.pkl`, `stylegan3-r-ffhqu-1024x1024.pkl`,`stylegan3-r-ffhqu-256x256.pkl`
>
> `stylegan3-t-metfaces-1024x1024.pkl`，`stylegan3-t-metfacesu-1024x1024.pkl`,`stylegan3-r-metfaces-1024x1024.pkl`, `stylegan3-r-metfacesu-1024x1024.pkl`,`stylegan3-t-afhqv2-512x512.pkl`,
>
> `stylegan3-r-afhqv2-512x512.pkl`

```shell
#使用预训练的AFHQv2模型生成图像
python gen_images.py --outdir=out --trunc=1 --seeds=2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

#使用预训练的AFHQv2模型生成视频
python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
```

上述命令的输出放在`out/*.png`下，由`--outdir`控制。`--network`可以指定url路径(如上例代码)或者本地`pkl`路径。

下载的网络pickle缓存在`$HOME/.cache/dnnlib`，可以通过设置`DNNLIB_CACHE_DIR`环境变量来覆盖。默认的PyTorch扩展构建目录是`$HOME/.cache/torch_extensions`，可以通过设置`TORCH_EXTENSIONS_DIR`来覆盖。



#### 训练与预训练

可以使用`train.py`训练新的网络，例如：

```shell
# 例:训练 StyleGAN3-T 数据集AFHQv2, 4张DCU.
python train.py --outdir=~/training-runs --cfg=stylegan3-t --data=~/datasets/afhqv2-512x512.zip \
    --gpus=4 --batch=16 --gamma=8.2 --mirror=1
  
#例:预训练的FFHQ-U pickle开始，使用1 DCU微调MetFaces-U的StyleGAN3-R
python train.py --outdir=~/training-runs --cfg=stylegan3-r --data=~/datasets/metfacesu-1024x1024.zip \
    --gpus=4 --batch=16 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 \
    --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.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`中记录各种统计数据。



### 性能和准确率数据



##### DCU训练损失

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



##### 生成效果（Metfaces）

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



##### 频谱分析

![image-20230710182732783](./pngs/acc.png)

### Note

`--gamma`: per-dataset basis R1 regularization weight，最需要调整的参数，通常和训练集的分辨率呈 *倍率平方* 关系，如分辨率从 256x256 → 512x512 那么对应的 gamma 则从 2→8。`-r` 和 `-t`是一致的，略低于 stylegan2。

`–-metrics`：用于在训练过程中评估生成的图像相较于数据集的质量，如果不是为了写paper做研究性数据就设置为none，否则**非常耗时**。

`--aug`:`noaug` 会禁用 ADA。在至少 100k 训练 images 时（包括翻转后的），这个效果不大，但是训练数据少于 30k 的时候，很有用。

`--cbase=16384`会加速训练，但是效果会变差（低分辨率几乎没有影响，如256*256）。



### 源码仓库及问题反馈

[https://github.com/NVlabs/stylegan3](https://github.com/NVlabs/stylegan3)



### 其他参考

[README_ORIGIN.md](./README_ORIGIN.md)

[Training configurations](https://github.com/NVlabs/stylegan3/blob/main/docs/configs.md)