代码收藏家技术教程 2022-07-20

深度学习模型试跑(十三):stylegan3

前言

在深度学习中，训练数据量不足常常会影响分类算法的性能。我从这几年的相关工作经验感受得出，缺乏训练数据并不是例外而是一种规律，这就是为什么很多人会想出各种各样的数据增强方法。吴恩达也说过，scale drives machine learning progress，也是对在深度学习领域数量影像质量这一概念的一种诠释。
我们可以使用常规的数据增强手段，比如参考链接中提到的使用的例如旋转翻转，旋转，裁剪，变形，缩放等各类操作原始数据来生成新的训练数据，但这并不能给我们带来真正的新图像。
相反的，当我们处理稀少的训练数据时，GAN数据增强方法就显得更具优势。假设我们想要训练在工业生产中某些特定的缺陷如刮痕、脏污、磨损等，但常常会遇到这类希望出现的缺陷很少产生的现象，这就导致了我们可能只有一小部分显示典型缺陷的图像来训练网络。如果我们使用GAN，我们就可以为任何给定的缺陷类型生成额外的“真实的”图像。
这篇我简单地记录一下在我自定义的数据集上训练最近大火的模型“StyleGAN3”，对GAN感兴趣的，可以到B站找相关的资料，例如李沐大神的GAN精读，以及唐宇迪大神的GAN系列解读等。

我的环境：

pytorch 1.10.0

CUDA 11.5，cudnn 8.2

VS2019

Pycharm

由于部分算子是用CUDA 11.1来写的，所以CUDA要求11.1+。目前手上算力不足的建议只简单测试一下或者暂时放弃，不然训模会很痛苦，这里给出StyleGAN2的训练指标，显卡使用的是NVIDIA Tesla V100，大家可以自行参考。
请添加图片描述

一.解读

~~可以参考StyleGAN3论文解读，模型关键的地方都有做了些许介绍。如果对StyleGAN3十分感兴趣~~ ，建议阅读StyleGAN3的论文。

解读部分更新于2022/3/22,还在持续更新中

1.1 论文解读

最近重新看了StyleGAN3的论文，核心部分应该在第三段Practical application to generator network，这一段具体介绍了StyleGAN3相较于StyleGAN2在生成器方面的改进方法，这里面涉及了很多数字信号处理（复变函数）的处理手段，于是论文在附录(Appendices)中花了将近20多页介绍了这方面相关的理论知识。总而言之，整篇文章都在表现出作者是在用数字信号的概念来诠释网络的构建流程，光这个相对比较新奇的思路就让很多人阅读整篇论文都感觉很晦涩。为了防止被很多人冲，我这里大部分借鉴了某位大佬的解读，里面加入自己些许的理解，如果大家有疑问欢迎在评论区友善地提出。

1.1.1 整体逻辑

建模：将特征图像看作由若干（无限）个离散的二维的网格组成，每个网格都看作是一个δ函数/单位脉冲信号(a值不同，但采样率为s)，这些信号通过网络的时被看作为连续信号，实际使用的特征图只是这些连续信号的离散样本，其实就是常用的时频转换。

问题分析：发现目前的网络架构没有一个明确的机制来限制生成器以严格的分层方式合成图像。所谓分层，可以认为是因果的，粗糙特征对细微特征具有控制效果，比如在人脸图像中，转动脸部会带动鼻子的转动，鼻子的转动会进一步带动鼻子上毛孔的转动此外，由于特征图的频率不满足奈奎斯特-香农（Nyquisit-Shannon）采样定理的条件，会出现混叠现象。

解决方式：重新设计一个无混叠（alias-free）并严格遵循分层合成方式的网络。

分析：设计出来的无混叠生成器在平移或旋转时都是等变的，且这种生成器不会产生纹理沾粘（texture sticking）的现象。

1.1.2 动机

连续变量与离散变量之间的相互转化
作者利用数字信号处理中的概念，并解释通过网络的信息流为空间无限连续信号。我们实际上使用的特征图只是目标画布中连续信号的离散样本。它们可以被看作是连续信号的一种方便的编码。如果我们将连续信号中的单位正方形[0,1]设为我们的目标画布，那么feature map的大小就可以表示连续信号转换为离散信号时的采样率。我们所说高频或低频的讨论的是对连续信号进行傅里叶变换后在频域中得到的频率。由于是采样，因此需要满足奈奎斯特-香农采样定理的条件。也就是说，连续信号的最高频率必须小于采样率的一半(这通常称为奈奎斯特频率)，否则就会出现混叠问题。关于这个定理的理解可以看这个知乎链接，具体作用可以看这份代码。

当前GAN网络的缺陷
作者发现目前流行的GAN结构都没有从一种自然的分层次方式（hierarchical manner ，由浅到深，由粗到细，由低到高频率合成特征）来合成图像，尽管这些GAN网络已经限制了各层特征图的分辨率，使得浅层的特征图不能代表高频信号，但不能保证各层操作产生的新频率小于对应的奈奎斯特频率。如果不满足了上述条件，就会出现混叠问题，使高频在频域中被表示为低频，污染整个信号。

主要贡献
作者设计一个网络体系结构，严格遵循理想的分层次方式。每一层都被限制在我们指定的范围内合成频率，因此，消除了混叠的问题。

1.1.2 方法

重新设计基础算子
很多Gan网络包含基本的操作，如Conv、上采样、下采样、非线性。下面，我们将分别对它们进行分析，看看它们是否存在混叠问题。如果是这样，我们该如何修复它们。
其中，
（1）下采样将信号重新采样到较低的采样率，即从高采样率s到低采样率s’。它是用来缩小频谱中的可行区域的，实际就可以看作是一种阈值化处理。值得注意的是，之后的采样率可能小于原始信号最高频率的两倍。因此，需要事先使用低通滤波器来限制原始信号的频率，使其小于降低的采样率的一半，然后才能进行下采样过程。

（2）上采样将信号重叠为更高的采样率。它用于在频谱中添加余量，让可行区域更大（使后续层可以引入新频率），但它本身不会引入新频率。该过程通过将原始信号与初始信号0交叠来实现该过程，然后使用低通滤波器来移除频域中的成像，这里低通滤波器使用的截断频率是s/2，采样率是s’。

（3）非线性函数/激活函数，在代码filtered_lrelu.py/_filtered_lrelu_ref()中具体使用了。基本函数非线性如ReLU等被视为用来引入新的频率，非线性引入的新频率包含两部分:第一部分满足采样定理的条件，第二部分不满足。我们要保留前者，消灭后者。然而，如果我们直接将非线性应用于离散特征映射，新引入的第二部分频率将直接产生混叠。因此，作者提出了一个非常有趣的方法:首先，用m对信号进行上采样(通常设为2)，然后应用非线性，最后对信号进行下采样( upsample-nonlinearity-downsample)。第一个上采样是增加奈奎斯特频率，为新引入的第二部分频率增加净空间，以避免混叠。然后，下采样过程(包括用低通滤波器以消除第二部分频率)将信号转换回其原始采样率。（论文D.1 Gradient computation）

（4）低通滤波器的设计基于Kaiser窗函数（论文C.1 Kaiser low-pass filters），这里可以看它的直观作用。

等变和纹理沾粘
等变是指当输入平移/旋转时，输出也平移/旋转。(stylegan3/metrics/equivariance.py )
（1）平移等变（代码 equivariance.py[line 224] Integer translation (EQ-T)）
根据上面的理论分析，如果我们将信号在整个网络的时域上视为无限连续信号，那么信号在时域上的平移实际上并不会改变信号在频域上的幅值。无论在时域上如何移动(上下左右)输入信号，网络每一层的输出都会跟着它移动，最终的输出信号肯定也会跟着它移动。作者定义了一个度量来计算平移等变方差:EQ-T（论文公式3）。作者反映出峰值信噪比 (PSNR)以两组图像之间的分贝 (dB) 为单位，通过将合成网络的输入和输出平移某个随机量获得。
（2）旋转等变（代码 equivariance.py[line 243] Rotation (EQ-R)）
对于旋转等变，需要对卷积和低通滤波器(LPF)做一些修改(论文 E.3 Rotation)。作者认为卷积核函数在时域是径向对称的，因为如果旋转输入信号，最直观和简单的方法是对Conv核执行同样的旋转，然而如果这样的话，两者之间就没有相对的运动，相当于原来的操作。对于低通滤波器的解释理论上和卷积同样的道理。EQ-R（论文公式23）。
（3）纹理沾粘
等变网络不存在这种现象。这一现象的表现是高频和低频特征不会以相同的速度同时变换。但如果网络具有等变性，那么所有的特征必须以相同的速度变换在一起，这种现象自然不会发生。（论文 Figure 1）

整体网络架构的详细设计
相较于stylegan2的生成器，除了基本操作的变化，网络架构也发生了变化，具体可参见论文Figure 3中带有**+号的Configuration，接下来会一一介绍。~~Digression：听着ABCDEFU，写着这冗长的博客，真的应景。~~
*(config B)*用傅里叶特征替换StyleGAN2中学习到的输入常数
根据前面的分析，我们处理的输入本质上是一个无限的连续信号，所以作者在这里使用了傅里叶特征，它具有空间无限的特征即离散输入信号可以从连续表达式中采样。同时，由于存在一个实际的连续表达式，我们也可以很容易地对信号进行平移和旋转，然后对其进行采样并输入到网络中，这样我们就可以方便地计算EQ-T和EQ-R，这里可以详见代码networks_stylegan3.py[line 230~234]**。

        x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) #完成通道展平后采样梯度叉乘转换频率 
        x = x + phases.unsqueeze(1).unsqueeze(2)
        x = torch.sin(x * (np.pi * 2))#用sin函数计算
        x = x * amplitudes.unsqueeze(1).unsqueeze(2)#信号衰减

*（config C）*删除了每个像素的噪声输入，因为它们与自然转换层次结构的目标非常不一致，也就是说，每个特征的确定的亚像素位置完全继承自底层的粗糙特征。因此，作者减少了mapping network的深度（代码 train.py[c.G_kwargs.mapping_kwargs.num_layers], 由8变为2）并且去掉了 mixing regularization和path length regularization（代码 train.py[line 233~251],stylegan2与stylegan3-r的区别，stylegan3-r就没有这两项）

*（config D）*在训练过程中跟踪所有像素和特征图上的指数移动平均值σ^ 2=E[x^2]，并将特征图用√σ ^2来划分（实际上使用卷积来划分以提高效率）,详见代码networks_stylegan3.py[line 153~155]；消除跳过连接，改为使用sigma的EMA归一化

*（config E）*边界和上采样，一种非常直观的方法（临界采样）是将低通滤波器的截断频率Fc设置为采样率S的一半即S/2，将过渡带f频率设为fh = (√2 − 1)(s/2)。详见代码networks_stylegan3.py[line 436]

（config F）滤波的非线性，使用当前深度学习框架中原始的构造来实现上样本👉leaky ReLU👉下样本这一序列其实并不有效，因此作者实现了一个自定义的CUDA内核(Appendix D)，它结合了这些操作(Figure 4b)，从而加快了10×的训练和相当大的内存节省。具体的代码可以看torch/unit/ops/filtered_lrelu开头的代码，尤其是cuda代码。

*（config G）*非关键采样，因为低通滤波器只是近似值，所以它不是频域中的理想矩形窗口，因此会有一些缺失的频率仍然可以通过临界点。为了抑制混叠，可以简单地将截止频率降低到fc=s/2−fh，从而确保所有混叠频率(在s/2以上)都在阻带内。由于信号现在包含的空间信息较少，作者修改了用于确定特征映射数量的触发方式，使其与fc成反比，而不是采样率s。networks_stylegan3.py [line 431]

*（config H）*可变的傅里叶特征，networks_stylegan3.py [line 204~215]，这里引入一个学习到的仿射层，它为输入的傅里叶特征输出全局平移和旋转参数，该层被初始化以执行身份转换，但会随着训练的时间推移而是使模型学习使用该机制。

*（config T）*灵活的层表述，虽然作者们发现他们对网络的改动已经大大提高了网路的等变性，但一些可见的伪影仍然存在，正如论文附带的视频所示。针对这一问题，作者们建议对每一层分别进行设计，他们希望在低分辨率层中有尽可能大的衰减，在高分辨率层中保留更多的高频特征。networks_stylegan3.py [class SynthesisNetwork]
在Figure 4c中展示了一个14层生成器中滤波器参数的示例，最后有两个严格采样的全分辨率层；
截止频率从第一层的fc=2几何增长到第一临界采样层的fc=sN/2；
最小可接受的阻带频率从f_{t,0} = 2·power(2.1) 开始，它虽然也是不停地增长但显然慢于截止频率。在作者的测试中，最后一层的阻带频率是ft=fc·2·power(0.3);
采样率s被设置为大于ft的2的最小倍数的两倍(但不超过最终的输出分辨率);
将过渡频率的半宽设置为fh = max(s/2, ft) -fc等

*（config R）*旋转等变性。1.作者在所有层上用1×1替换3×3卷积，并通过增加两倍的特征映射数量来弥补减少的容量；networks_stylegan3.py [line 296]。 2.作者用一个径向对称的基于jinc的滤波器来代替基于sinc的降采样过滤器，我们使用相同的凯撒方案构造了该滤波器networks_stylegan3.py [line 376~384]

二.模型搭建(以win 10为例)

参考，首先git clone https://github.com/NVlabs/stylegan3.git克隆到本地，由于项目没有写requirements.txt，我把我环境中相关库导出在这里以供参考。

absl-py==1.0.0
addict==2.4.0
# Editable install with no version control (basicsr==1.3.4.8)
-e c:\programdata\miniconda3\envs\gan\lib\site-packages
cachetools==4.2.4
certifi==2021.10.8
charset-normalizer==2.0.8
click==8.0.3
colorama==0.4.4
cycler==0.11.0
fonttools==4.28.2
future==0.18.2
glfw==2.4.0
google-auth==2.3.3
google-auth-oauthlib==0.4.6
grpcio==1.42.0
h5py==3.6.0
hdf5storage==0.1.18
idna==3.3
imageio==2.12.0
imgui==1.4.1
importlib-metadata==4.8.2
kiwisolver==1.3.2
lmdb==1.2.1
Markdown==3.3.6
matplotlib==3.5.0
networkx==2.6.3
ninja==1.10.2.3
numpy==1.21.4
oauthlib==3.1.1
opencv-python==4.5.4.60
packaging==21.3
Pillow==8.4.0
protobuf==3.19.1
psutil==5.8.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pyparsing==3.0.6
python-dateutil==2.8.2
PyWavelets==1.2.0
PyYAML==6.0
requests==2.26.0
requests-oauthlib==1.3.0
rsa==4.8
scikit-image==0.18.3
scipy==1.7.3
setuptools-scm==6.3.2
six==1.16.0
tb-nightly==2.8.0a20211202
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.0
tifffile==2021.11.2
tomli==1.2.2
torch==1.10.0+cu113
torchaudio==0.10.0+cu113
torchvision==0.11.1+cu113
tqdm==4.62.3
typing_extensions==4.0.0
urllib3==1.26.7
Werkzeug==2.0.2
wincertstore==0.2
yapf==0.31.0
zipp==3.6.0

记住把"C:\Program Files (x86)\Microsoft Visual Studio<你的版本>\Community\VC\Auxiliary\Build\vcvars64.bat"添加到path环境变量里。
请添加图片描述
在torch_utils/ops/filtered_lrelu.py中第32行添加extra_cflags=['/std:c++17'],。

至此模型搭建基本完成，可以用如下命令进行单个图片/视频的生成简单体验一下。如果报错缺库，就自行pip安装。

# Generate an image using pre-trained AFHQv2 model ("Ours" in Figure 1, left).
python gen_images.py --outdir=out --trunc=1 --seeds=1 \
    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

# Render a 4x2 grid of interpolations for seeds 0 through 31.
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

请添加图片描述

二.模型训练

可参考教程，这里只介绍我如何训练我自己的数据。

制作自定义数据:

yolov5训练的流程图，只为了说明我的目的是发现图中类似于0的缺陷。

准备一组特征相似的图片，放在同一个文件下，并将它们的尺寸全部转换成256*256。这里我考虑到了训模的时间，如果资源充足可以不进行缩放。

import os
import cv2
import sys
import numpy as np

path1 = r'D:\lbq\dataset\PG\2021_12'
path2 = r'D:\lbq\dataset\PG\dent'

for filename in os.listdir(path1):
    if os.path.splitext(filename)[1] == '.bmp':
        # print(filename)
        sas = os.path.join(path1, filename)
        img = cv2.imread(sas)
        tem = cv2.resize(img, (256, 256))
        print(filename.replace(".bmp", ".jpg"))
        newfilename = filename.replace(".bmp", ".jpg")
        # cv2.imshow("Image",img)
        # cv2.waitKey(0)
        dst = os.path.join(path2, newfilename)
        cv2.imwrite(dst, tem)

数据打包:

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

python dataset_tool.py --source=数据集路径 --dest=转换后数据包.zip

训练:
这里我说下需要着重考虑的参数，我参考了官方训练配置指南和一个StyleGAN2的训练教程。

–cfg：有StyleGAN3-T (等变平移)、 StyleGAN3-R (等变平移+旋转)、StyleGAN2这三项，优先选择前两项。

–batch：总批次大小，需要根据gpu的配置情况而选择合适的值。我这里只有一个3090，试了几次发现最高只能到24。

–gamma：R1正则化权重，根据作者的解释，值越大模型越稳定，值越小模型多样性越强。这个跟IS评价指标类似，涉及了一个熵值的问题。我这里选了参考StyleGAN2，选了10这个比较大的值。

–kimg：类似于iterations，一个总的迭代次数。默认值是25000，但作者说5000基本上效果就很好了。

–tick与–snap与：前者表示间隔多久打印一次训模信息，后者表示在tick*snap的基础上多久保存一个模型以及该模型的一张推理的结果。

–workers：windows设为0或1，懂得都懂。。

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

#开始训练：
python train.py --outdir=training-runs --cfg=stylegan3-r --data=datasets/dent.zip --gpus=1 --batch=24 --gamma=10 --mirror=1 --kimg=1000 --snap=50 --workers=1 --batch-gpu=12 --metrics=none
#恢复训练：
python train.py --outdir=training-runs --cfg=stylegan3-t --data=datasets/ffhq_biked.zip --gpus=1 --batch=16 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 --workers=1 --batch-gpu=16 --metrics=none --resume="D:\lbq\code\stylegan3\training-runs\00005-stylegan3-t-ffhq_biked-gpus1-batch16-gamma6.6\network-snapshot-000200.pkl"

第800个kimg的预测结果，花了整整一天的时间！
请添加图片描述

判别器的scores/fake和scores/real这两个指标曲线看上去都还不错。

图像生成:

借助搭建模型里的预测命令进行生成，注意seed配置。

python gen_video.py --output=dent.mp4 --trunc=1 --seeds=0-31 --grid=4x2 --network=training-runs/00000-stylegan3-r-dent-gpus1-batch24-gamma10/network-snapshot-001000.pkl

下面是预测的视频，我转换为jif图片。CSDN对插图上传大小有要求，于是我就压缩了该图像，导致了些许模糊。
请添加图片描述

三.问题汇总

command ‘[‘ninja’, ‘-v’]’ returned non-zero exit status 1
把环境中的site-packages/torch/utils/cpp_extension.py
[‘ninja’, ‘-v’]
改成
[‘ninja’, ‘–version’]
ImportError: DLL load failed while importing bias_act_plugin
参考
权重文件下载
链接
AttributeError: module ‘distutils’ has no attribute ‘version’
链接
RuntimeError: The size of tensor a (256) must match the size of tensor b (128) at non-singleton dimension 1
后面加了一个–cbase=16384 解决
模型
链接：https://pan.baidu.com/s/12-VlDG20AkCtUvVJQnKhyg
提取码：xzt0