README_old.md 6.34 KB
Newer Older
mashun1's avatar
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
## [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)