源代码地址:https://github.com/NVlabs/stylegan2-ada-pytorch文章来源地址https://www.toymoban.com/news/detail-409469.html
- 这是一篇代码阅读笔记,顾名思义是对代码进行阅读,讲解的笔记。对象是styleGAN2的pytorch版本的代码,在github上有一个开源库。一边笔记方便我回顾,一边也对深度学习初学者有一些阅读理解代码的示例作用吧。代码毕竟是基本功,看到我发的一些代码,评论区问的一些问题实在是代码基础不行。
- 读一个代码的目的很多样,有些只是想看它的预处理,有些只是想看他的模型,不同目的阅读方式和详略程度也不同。我这里是为了全部读懂给女朋友讲解的,所以会以全部读懂为目的来进行阅读和笔记。
train.py
- 起始点有很多种,我习惯从主函数开始读。train.py是主程序文件,直接拉到最后可以看到,主函数是从main函数开始的,那就找到main函数,438行开始
train.py/main()
- 482行看名字就知道是输出日志用的,不重要,可以先跳过
- 486行,可以看到一些主要参数是在这里设置的,后续要找再来看,这里先跳过
- 491-498行,设置了一些输出路径,也是跳过,有需要再来看
- 520-524行,将训练设置保存到一个json文件中
- 526行到533行是程序的主体部分,在这里启用了多线程,运行
subprocess_fn
函数,因此下一步就看这个函数。这里稍微展开说一下这个多线程是怎么回事,就是利用了torch.multiprocessing
实现了每个GPU分配一个线程,并且多线程之间是用spawn方式创建的。也就是说,你有多少个GPU,就会同时运行多少个subprocess_fn
函数,并且spawn方式意味着这些线程都有独立的python解释器程序,资源是复制的,有自己的独立内存而非全部共享内存,而529行是指定了一个临时路径用来给这些线程进行交流,在这个路径下实现需要共享的部分变量。
train.py/subprocess_fn()
- 367行到380行是torch分布式训练的一些初始化设置
- 主程序在training_loop模块的training_loop方法中,接下来就跳到这里
training/training_loop.py
training/training_loop.py/training_loop()
- 这里终于遇到第一个关键点,136行,数据集,调用了
construct_class_by_name
函数。后面会讲解construct_class_by_name
这个函数,这里只需要知道它是个根据输入的参数,返回一个根据参数确定的类的方法即可。最近越来越多的深度学习代码使用这种包装方式,本质上就算想用字符串来调用类,又为了代码统一和简洁,包装得一层接一层,读起来是真的麻烦,而且类名隐藏起来了,甚至无法用vscode的智能追踪来找这里用到的到底是什么类。 - 这里直接说,training_set根据train.py的107行,是training.dataset.ImageFolderDataset类的对象
- 以及150行和151行,
G
和D
根据train.py的176 177行分别是training.networks.Generator类和training.networks.Discriminator类的对象。而G_ema
是G的一个指数移动平均版本,在训练过程中,G
的参数会随着step而更新,而G_ema
是G
的迭代过程中各个时期的参数的指数移动平均版本,相比G
,G_ema
的变化更加柔和,这是个常用的技巧。 - 154到159行的代码加载了模型的参数,可能是预训练的也可能是训练中止接着跑的。
- 175行
augment_pipe
是train.py 287行 training.augment.AugmentPipe类的对象 - 180到190进行了多线程训练的模型包装
- 192到214行定义了训练的几个阶段。这里展开解释一下,GAN的训练策略相比普通模型稍微有些复杂,训练是分阶段的,每个iteration通常要分别训练G和D,并且在训练G的时候,D的参数要固定,训练D的时候,G的参数要固定。这段代码定义了4个阶段:Gmain Greg Dmain Dreg。
- 195行 loss 根据train.py 187行,是training.loss.StyleGAN2Loss类的对象
- 199行 opt 根据train.py 185行,是torch.optim.Adam类的对象
- 216-227行从训练集中采样了一些图片进行可视化(可做debug用),同时也将还没训练的G的输出也做了可视化(可用于检查resume是否加载或者pretrain模型的初始性能)
- 259行开始训练
- 260行获取真实图像和对应的label,261-262行归一化图像并划分图像和label到各个GPU
- 263行-264行生成随机向量作为Generator的输入,并划分到各个GPU
- 265-267行并从训练集中随机采样label作为条件label,并划分到各个GPU
- 270行,依次迭代前面提到的4个阶段
- 278行,把这个阶段需要训练的module设为计算梯度(如训练G的时候,设G的requires_grad为True,而D的为Flase)
- 281-284行,根据当前所处阶段,为每个GPU分别计算损失。每个阶段的损失介绍损失的时候会展开。
- 287-294行,根据当前所处阶段,更新待训练的参数(如Gmain和Greg阶段就只更新G的参数),并且把之前设为True的requires_grad改回去,然后进入下一阶段,直到4个阶段全部完成。
- 296-305行,为G计算指数移动平均,从而更新G_ema的参数
- 311-315行,根据训练过程的损失,调整数据增强策略的参数,具体在下面介绍数据增强的时候会展开。
- 318-320行,这里是设置了continue条件,使得每迭代4000(
kimg_per_tick*1000
)张图片才会运行322行以后的内容一次。实现方式是cur_nimg
会一直增加,而tick_start_nimg
只有在下面的代码会被设置为cur_nimg
,这样一旦运行了一次下面的代码,下次判断小于号就会成立,直到cur_nimg
增加了4000使得小于号不成立,然后又会运行一次下面的代码。而done条件是因为,break出循环之前需要运行一次下面的代码,所以设置了当迭代图像数满足图像总数的1000倍的时候,就要退出了,这时候不管是不是每4000次的间隔到了,我都要往下走。 - 341行设置了另一种退出的方法(代码里似乎没有设
abort_fn
所以应该这一段代码是没有用到),可以为training_loop
传一个有效的abort_fn
,使得如果准确率等满足条件返回True,从而不需要跑满1000epoch可以退出。 - 348-350行为当前iteration生成的图片保存到本地,因为这段代码在322行之后,所以每迭代4000张图片才会生成一次。
- 353-367行保存了模型参数,同理也是4000张图片才保存一次。
- 370-379行计算了指标,后续会展开介绍
- 381-389行和322-338行都说计算运行时间和存储消耗的,就跳过了
- 391-406行都是输出日志的,跳过
- 414行是while True循环的唯一退出点。如果运行完成了,就从这里退出循环,结束训练。
dnnlib/util.py/construct_class_by_name()
- 这个函数只有两行,调用了
call_func_by_name
函数并以其返回值作为自身的返回值。call_func_by_name
函数定义在279行,调用了get_obj_by_name
函数,并进一步调用得到的func_obj
,以func_obj
的返回值作为call_func_by_name
的返回值。所以这里其实就是调用了get_obj_by_name
函数得到了类,func_obj保存的就是得到的类,然后实例化并返回,所以返回的是类的实例化对象。 -
get_obj_by_name
函数在273行,调用了get_module_from_obj_name
和get_obj_from_module
。有点绕,其实是因为,name是xx.yy.zz的格式,zz才是类名,xx.yy是模块名,所以先调用222行的get_module_from_obj_name
从xx.yy.zz中提取出xx.yy和zz,然后再借助get_obj_from_module
函数从xx.yy模块中调用zz类。 -
get_module_from_obj_name
函数的核心就在231-239行,231-232行其实就是给出根据“.”的位置对字符串划分成两部分的全部可能,所以如果是xx.yy.zz就会被拆成xx和yy.zz或者xx.yy和zz。然后在235到239行,对每种可能性都进行尝试,尝试从xx.yy中import zz,尝试从xx中import yy.zz,因为用的是try,试不出来可以继续,直到试出来,就知道正确的划分方法是什么。 - 而
get_obj_from_module
函数是通过269行的getattr函数来获取模块中的类的。
training/dataset.py
-
ImageFolderDataset类定义在training/dataset.py的154行,是同文件24行Dataset类的子类,一般看__getitem__函数即可。返回值有
image
和label
。image
根据210-220行的重写,是一个CHW的unit8(0-255)的np array。label
是onehot的float32的np array
training/networks.py
-
Generator类定义在training/networks.py的477行。476行是一个装饰器,意思是调用
training.networks.Generator
的时候,实质上返回的是persistence.persistent_class(Generator)
,这个装饰器只是为这个类添加了一些辅助功能,不影响接下来的理解,所以先跳过,后续会解释这个装饰器,先接着看模型 - 模型由两个子模块组成:MappingNetwork 和 SynthesisNetwork
training/networks.py/MappingNetwork()
- MappingNetwork定义在174行,从初始化函数看起,200行前面定义了一些变量的维度,201行定义了中间全连接层的维度。
- 204行定义了第一个全连接层,当使用condition label的时候,对这个one hot的condition label进行embed,embed后的特征将和z连在一起作为后续网络的输入。
- 205-209行定义了网络的主体全连接层
- 211-212行定义了一个名为
w_avg
的变量,它不会随着step更新值,但会在一些特殊的时刻进行值的更新和被使用。 - 这里的FullyConnectedLayer(89行)相比普通的全连接层的区别在于,当
lr_multiplier
不为1时(208行定义的就不为1,是0.01),这些层的参数的学习率和其它参数的学习率相比会乘以一个lr_multiplier
(具体实现其实就是把参数直接乘以一个lr_multiplier再去用,实际效果就等同于学习率乘了一个倍数,因为计算这些参数的梯度的时候也是会因此乘以一个lr_multiplier
导致step的时候步长会乘以一个lr_multiplier
的) - 接下来看forward函数。219和222行都仅仅是检查向量的shape。
normalize_2nd_moment
函数看21行,其实就是先统计这些特征值的标准差(每个样本单独统计),接着除以标准差进行归一化。其实这么说不太准确,因为没有减去均值,仅仅是先平方,然后平均,然后开根,然后除(rsqrt是1/sqrt)。而20行的装饰器仅仅是使得torch.autograd.profiler.record_function
能跟踪到这个函数而已。至于torch.autograd.profiler
后续会介绍是个什么东西。 - 然后在223-224行,c向量送进一个全连接层编码,归一化,然后和归一化后的z向量被concatenate到一起,作为后面全连接层的输入
- 226-229行就是主体的mapping network,对合并的编码和z向量前向传播经过几层全连接层
- 231-234行保存全连接网络的输出的移动平均(
lerp
是根据w_avg_beta
对w_avg
和x
进行插值的函数)到w_avg
变量中 - 236-239行重复了
num_ws
份x
,放在dimension 1上,也就是说现在shape是(B,num_ws,w_dim)
,具体num_ws
是什么下面介绍SynthesisNetwork时会展开说明 - 242-248行,查完整份代码没有看到哪里有把
truncation_psi
设为非1的值,所以理论上正常情况这部分代码是不会运行到的。看意思应该是利用w_avg
对x
进行进一步移动平均,这里的移动平均就是对x做了,影响的是x
的值,前面的移动平均只是存下来而已,对实际训练过程不会有什么影响。之所以说是截断,是因为当x
在训练过程中突然出现异常大或者异常小的值时,这段代码可以通过移动平均限制这些值不要偏离正常范围太远。
training/networks.py/SynthesisNetwork()
- SynthesisNetwork定义在424行。首先看init函数,440行根据要生成的图片的分辨率,定义了各个block的resolution,依次是2的2,3,4,。。n次方,使得2的n次方最接近要生成的图片的分辨率。441行则定义了各个block的通道数为32768除以block的resolution,但最小是512。
- 442定义了一个称为
fp16_resolution
的变量。FP16是一个降低运算量和内存占用的技巧,将32位浮点运算用半精度运算来近似。模型对分辨率最高的num_fp16_res个block进行FP16计算,所以这里是在算开始进行FP16计算的block的resolution。在448行当block的resolution大于等于这里算出来的fp16_resolution
时,意味着这个block要进行FP16计算而非全精度的计算。 - 445-455行定义了SynthesisNetwork的主体由堆叠的几个SynthesisBlock组成。这里还统计了num_ws,后续会解释这个是什么。这里展开解释一下455行的
setattr
函数,是一种通过字符串变量定义类成员名的方法,比如setattr(a,'hah',1)
,那么当调用a.hah
的时候,返回值会是1,也可以用464行的getattr
函数实现调用。 -
SynthesisBlock后面解释,先接着看forward函数,forward函数的输入是MappingNetwork的输出,即是全连接并repeat了
num_ws
遍后的编码特征,shape为(B,num_ws,w_dim)
,也就是说对于每个batch,有num_ws
个重复的w_dim
维的特征。为什么要重复,我的理解是这些副本在后续会被各个模块分别使用,可能是为了避免相互影响? - 463-466行将输入的ws特征在dimension 1上拆成多份,每一份分给一个block。也就是说现在每个block的输入是(B,num_conv,w_dim),其实就是每个block分到了num_conv个重复的特征向量。最后一个block会得到 num_conv+num_torgb份(因为只有最后一个block的num_torgb不为0)
- 468-471行则开始前向传播,每个block的输入是分得的ws和上一个block的输出(x和img),第一个block输入的(x和img)为None。最后一个block输出的img为SynthesisNetwork最终的输出
training/networks.py/SynthesisBlock()
- SynthesisBlock定义在329行,首先还是看init。354行定义了一个变量,和之前说的一样这个变量是不会被step更新值的。这里用到了upfirdn2d.setup_filter这个方法,参数是resample_filter,为[1,3,3,1],其实是定义了一个torch的tensor,是输入的外积归一化后的结果。也就是[1,3,3,1]和自身的外积,得到一个(4,4)的tensor,并且进行归一化使得元素和为1。所以得到的是一个2D的低通滤波器。
- 359行定义了一个随机变量,这个变量仅在SynthesisNetwork的第一个block被定义,作为起始的随机变量用来生成图片。
- 361-364行定义了第一个SynthesisLayer,这个类后续会介绍。这里定义的这一层,第一个block是没有的。
- 366-368行定义了第二层SynthesisLayer,这一层每个block都有
- 370-373行根据现有代码是一定会运行的(
architecture
就是’skip’,除了cfg
定义为’cifar’时,cfg默认是’auto’)所以每个block都定义了一个ToRGBLayer,后续会介绍。所以到这里可以看出,除了第一个SynthesisBlock为一个SynthesisLayer加一个ToRGBLayer外,其它的SynthesisBlock为两个SynthesisLayer加一个ToRGBLayer。 - 375-377行代码根据现有architecture的是不会运行的,不介绍。
- 接着看forward,381行,根据前面,ws是重复的特征向量,所以这里
unbind
就是把重复的那个维度解出来,再套上iter
变成迭代器,那么每次next(w_iter)
都会生成一个特征向量,并且每次next
生成的特征向量不是同一个,但是内容相同。这个特征向量其实就是MappingNetwork生成的编码特征向量。 - 382行这里,和前面提到的FP16呼应了。因为styleGAN随着tensor往后传递,生成的tensor是越来越大的,分辨率越来越高,为了节约显存同时提高运算速度,可以把后面几个block的数据类型从float32改成float16,这样节约了一半的显存。
- 383行定义了向量在内存中的存储格式,连续的存储可以提高某些运算的速度,同时有些运算要求连续存储的tensor才合法,但将存储整理为连续存储会消耗时间。
- 384-386行定义了
fused_modconv
bool变量,后续用到的时候再展开,这里只需要知道只有测试的时候才有可能是true。 - 389-391行,因为第一个block的
x
是None,所以这里就根据init中自己生成的随机变量来定义x
。从这个实现方式可以看出来一个关键信息,即只要模型定义了,随着训练过程,每次迭代,第一个block的输入x
都是固定的,不会改变。并且由于const
是torch.nn.Parameter,所以加载resume和load 已训练好的参数的时候也是生成和之前训练的时候同一个x。 - 396-406行即是主体SynthesisLayer的运行,可以看到SynthesisLayer是用来生成x的,而img则是ToRGBLayer生成的。
- 408-419行是生成img的过程。首先上采样上一个block生成的img(具体方式后续会介绍)。然后ToRGBLayer根据本block的x和ws,生成了残差图y,加到上采样了的img上形成本block输出的img
- 可以看出,其实网络中
x
的传递和img
是无关的,每个x
都只需要根据前一个block的x(第一个block则根据一个固定的x)和MappingNetwork生成的ws
,即可生成本block的输出x。而img则是根据前一个block的img
和本block的x以及MappingNetwork生成的ws
来生成的。
training/networks.py/SynthesisLayer()
- 这个类是SynthesisLayer的核心组成部分之一,定义在254行,先看init函数,274-284行定义了几个参数,分别是
resample_filter
,affine
,weight
,noise_const
(默认不会用到),noise_strength
(定义为0)和bias
。注意use_noise
在这份代码中默认是True的,所以if是一定会执行的。这几个参数的类型前面都介绍过,而具体作用看forward - 首先是290行,所以
affine
就是个把输入的w变成styles的全连接层。这里的w其实就是SynthesisBlock分给这个SynthesisLayer的ws
的其中一条特征向量。 - 然后是293行,代码默认
noise_mode
就是’random‘而且use_noise
是True,所以296不会运行,就是294行,重新生成了一个随机的噪声,满足0均值noise_strength
标准差的正态分布。值得注意的是,虽然noise_strength
在init中定义为0了,所以这里乘出来的noise
最初是一个全0的tensor,但noise_strength
是可训练参数,会随着训练过程变化,导致noise
后面还是会不为0的。 - 然后298行,up变量在每个SynthesisBlock(除了第一个block)的第一个SynthesisLayer被设为了2,其余都是1,所以flip_weight在每个SynthesisBlock的最后一个SynthesisLayer是True的,其余时候都是False的。
- 299行是主体,但是包装在了
modulated_conv2d
函数中,可以先到下面看完再回来这里看。 - 欢迎回来,302行的
gain
就是1,self.act_gain
根据276行是bias_act.activation_funcs['lrelu'].def_gain
,再看torch_utils/ops/bias_act.py的26行,是0.2,所以act_gain
就是0.2 - 303行
self.conv_clamp
根据train.py182行是256 - 304行可以先当作是
x
先加上self.bias
再进一个lrelu的激活函数,具体后面会展开,到此SynthesisLayer介绍完成,可以到ToRGBLayer
training/networks.py/modulated_conv2d()
- 这个函数定义在了同文件的第27行。47-49行为了防止FP16计算溢出,对
w
和styles
向量都除以了其各自的无穷范数(元素最大值)进行归一化,同时w
还除以了维度。demodulate
在SynthesisLayer中都没有设定,所以SynthesisLayer的modulated_conv2d
的demodulate
参数都是True。也就是说54到60行都会运行。 - 55-56行用到了
weight
和styles
两个tensor来构建w。weight
是在SynthesisLayer定义的shape为[out_channels, in_channels, kernel_size, kernel_size]
的卷积核,styles
是affine
这个全连接层输出的shape为[batch_size, in_channels]
的tensor。所以55行把weight
扩展了batch size那一维,变成了[1, out_channels, in_channels, kernel_size, kernel_size]
的tensor,然后styles
reshape成[batch_size, 1, in_channels, 1, 1]
的tensor,这两个tensor相乘,得到的是[batch_size, out_channels, in_channels, kernel_size, kernel_size]
,广播操作在这里的作用其实就相当于,把两个tensor都repeat成[batch_size, out_channels, in_channels, kernel_size, kernel_size]
,然后element-wise地相乘。直观上理解,这个相乘的作用是两个,一个是为同一个batch的不同sample分配不同的kernel,另一个是根据styles对每个channels的kernel进行rescale。 - 58行计算得到一个
dcoefs
,是w
的二范数的倒数,对第2 3 4维度分别算的,所以dcoefs
的维度是[batch_size, out_channels]
- 60行又用
dcoefs
来归一化w
,但这段代码和62-72行的代码只有一个会运行。当fused_modconv
为True时就运行60行的代码,否则运行62-72行的代码。这里用到了前面跳过的fused_modconv
,在386行,这个bool值只有在测试的时候,并且是在前面的FP32层才是True,在FP16层则必须当batch size为1时才为True。所以如果单看训练阶段,60行的代码是不会运行的,只有62-72行会运行。 - 64行,
x
是[batch_size, in_channels, H, W]
的,乘以一个[batch_size, in_channels]
的向量,会自动广播,其实就是把styles
repeat成x
的形状,再乘以x
- 65行调用了
conv2d_resample.conv2d_resample
来实现weight
对x
的卷积,具体后面会展开说,这里就先暂且当作普通的卷积。 - 再然后是66-71行的三个判断,SynthesisLayer中调用的
modulated_conv2d
,noise
都是tensor,不是None,所以只会运行67行,fma.fma
是自定义的一个函数,其实就是x
乘以dcoefs
加noise
。之所以这么写,而没有直接x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1) + noise.to(x.dtype)
,是因为可以利用torch.addcmul
来加速。 - 74到84行的代码也是在训练阶段是不会运行的,只有
fused_modconv
为True才会运行。相比63-72行差别在于卷积核从weight
改成了w
,并且因为乘以了styles
,每个样本有一个单独的卷积核,卷积结果也不需要乘以dcoefs
(应该是因为w
本身就是weight
乘以dcoefs
的缘故) - 所以其实到这里,
fused_modconv
的作用就很明显了。如果fused_modconv
是True,那么进行卷积的x
和w
,x
不做操作,w
是weight
乘以styles
再乘以dcoefs
的结果,并且对每个样本有一个卷积核;如果fused_modconv
是False,那么进行的是普通卷积,卷积的双方是x
和weight
,x
要乘以styles
,weight
不做操作,但卷积结果要乘以dcoefs
再输出。 - 看完这个可以回到SynthesisLayer继续下去
training/networks.py/ToRGBLayer
- 这个类和SynthesisLayer很类似,而且大多数函数都介绍过了,就不再一一介绍了。只需要注意一点的是,ToRGBLayer没有
Noise
,所以它调用的modulated_conv2d
中不需要加上noise
training/networks.py/Discriminator
- 这个类定义在673行。709行定义了Discriminator的主体由数个DiscriminatorBlock组成,同时还有一个MappingNetwork和一个DiscriminatorEpilogue。注意,711行定义的变量名,最小是b8(从693行可知),所以和715行的b4并不会冲突。还有就是,和SynthesisBlock相反,前面的DiscriminatorBlock分辨率更大,也就是说变量名b后面的数字更大,但排得更前。
- 接着看forward函数。可以看到Discriminator的输入由两部分组成,一个是
img
,一个是条件向量c
,根据719-721行,首先是按顺序调用堆叠的DiscriminatorBlock对输入的img
进行处理,然后调用一个MappingNetwork对输入的c
进行处理,最后用一个DiscriminatorEpilogue以DiscriminatorBlock的输出和MappingNetwork的输出作为输入,产生Discriminator最后的输出。
training/networks.py/DiscriminatorBlock
- 这个类定义在505行。和Genenrator不同,Discriminator的
architecture
是'resnet'
(只有cfg
为'cifar'
的时候为'orig'
)。而其它的,如FP16的设定和resolution都和Generator类似,就不细说了。 - 533-540行定义了一个迭代器,每次
next(trainable_iter)
都会返回一个bool值,用来判断当前层是否freeze,freeze多少层由freeze_layers
决定,默认参数下是0,也就是没有层会被freeze。如果要设定freeze多少层,可以通过train.py的freezed
参数设定,是一个int,指向从Discriminator的第一个block的第一个conv开始的全局层序数,也就是说如果freezed
设为5,那么从Discriminator的第一个block的第一个conv开始数,前5个conv都要freeze,后面的全都可训练。 - 542行的
in_channels
会在第一个DiscriminatorBlock为0,所以第一个block(分辨率最大的那个)是会运行543-544行的,而后面的block则不会。所以第一个block有一个额外的Conv2dLayer - 546-554行定义了3个Conv2dLayer,所以除了第一个block有4个Conv2dLayer外,其它的block都由3个Conv2dLayer组成
- 接着看forward,566-571行是第一个block独有的代码,其它的block不会运行这一段。这时
x
就是None,所以这段代码就是把输入的img
经过一个Conv2dLayer产生x
,同时把img
设为None,后续的block再也不需要用到img
了 - 575到578行就是把
x
分了两个支路,一个经过一层Conv2dLayer(在这个类内部会进行一次下采样使得分辨率变为原来的二分之一),一个经过两层Conv2dLayer(在第二层内部会进行一次下采样使得分辨率变为原来的二分之一),得到的两个支路的结果相加可以得到DiscriminatorBlock的输出。
training/networks.py/Conv2dLayer
- 这个类定义在123行,可以看作是一个卷积+上/下采样+激活函数。核心代码在
conv2d_resample.conv2d_resample
函数中,其它没什么难点,就偷懒一下不展开了。
training/networks.py/DiscriminatorEpilogue
- 这个类定义在615行,它的作用是以MappingNetwork的输出
cmap
和最后一个DiscriminatorBlock的输出x
作为输入,产生最终的分类值。637-640定义了4个layer,具体作用看forward函数 -
x
依次经过一个MinibatchStdLayer和一个Conv2dLayer,然后展平,送进两个全连接层,得到的结果和cmap
进行element-wise的相乘,然后全部求和并归一化得到最终的输出。这个输出同时也是Discriminator
的输出,是一个[batch_size, 1]
的tensor,就是对图片进行二分类的逻辑值,越高表示越real,越低表示越fake。
training/networks.py/MinibatchStdLayer
- 这个类定义在589行,不包含任何参数,仅仅是为
x
增加一些通道,这里初始化的group_size
是4,num_channels
是1 - 602行把
x
reshape成[4, N/4, 1, c, h, w]
,即把样本分成了N/4组,每组4个样本,然后对特征的各个维度分别计算组内标准差,再对每个位置每个通道的标准差取平均,得到[N/4, 1]
的向量,代表了每组的标准差,然后repeat到各个空间位置和组内样本上成为[N,1,H,W]
的向量,concatenate到x上成为新的一个通道,所以x变成了[N,C+1,H,W]
输出出去。
training/loss.py
training/loss.py/StyleGAN2Loss
- 整个py文件就这一个类的定义。24行,init函数传进来的参数,根据training_loop.py的184-190行,
G_mapping
是Generator的MappingNetwork成员,G_synthesis
是Generator的SynthesisNetwork成员,D
是Discriminator。注意,这个Loss函数是有成员的,有一个初始化为0的pl_mean
变量,在后续的accumulate_gradients
函数中会对这个向量进行移动平均。 - 接着看
run_G
函数,这个函数在accumulate_gradients
中被调用。41行的style_mixing_prob
是0.9(只有cfg
为'cifar'
时才为0)。 - 43行声明了一个满足均匀分布的随机整数,范围为1到
ws.shape[1]
.这里的ws
是Generator的MappingNetwork的输出。 - 44行有点复杂,先看where函数内的第一个参数,这是个随机的bool值,有0.9的几率为True,0.1的几率为Flase;第二个参数是刚才提到的1到
ws.shape[1]
之间的随机整数(均匀分布);第三个参数就是ws.shape[1]
。然后看where
函数,where
函数的三个参数按顺序依次是condition、input、other,意思是,如果condition是True,那么where
函数的输出就是input,如果condition是False,那么where
函数的输出就是other。所以这一行的意思是,cutoff
有0.9的几率会被置为1到ws.shape[1]
之间的随机整数,有0.1的几率会被置为ws.shape[1]
- 45行就是把
ws
在cutoff
后的那些向量替换成另一个根据随机z
和同一c
重新生成的向量。这里要回想一下ws
的生成过程,其实是一堆重复的特征向量堆叠而成的,也就是说原本ws[:, i]
和ws[:, j]
都是相同的特征向量。 - 所以40-45行的意思是,根据
z
生成的ws
,有0.9的几率会把其中随机数量的w
替换成另一个随机向量z2生成的w
,所以此时ws
内就有两种w
,一种是z
生成的w
,一种是z2生成的w
,并且比例是随机的。 - 47行根据调整后的
ws
调用Generator的SynthesisNetwork生成图片并返回。 -
run_D
很简单,就是先对图片augment,然后调用Discriminator判断图像的真假,产生逻辑值。 - 接着看
accumulate_gradients
函数,这个函数在每个阶段都会调用一次,所以进来的时候可能是四个阶段的其中一个。 - 如果是Gmain阶段,则损失函数是71行,对Discriminator预测的logit值取反算损失,softplus是个单调增函数,是平滑的relu,具体见(https://pytorch.org/docs/master/generated/torch.nn.Softplus.html#torch.nn.Softplus),所以backward会使得logit值增加,从而Generator生成的图像更真实,但此时Discriminator的参数要禁用require_grad,不参与更新,这一点在train_loop.py中进行了。
- 74行gain在这一阶段是1
- 如果是Greg阶段,79-80行,
pl_batch_shrink
是2,所以是把batch_size变成了原来的二分之一。这段代码的意义在于它能够使得Greg阶段的batchsize比Gmain阶段的batchsize小。 - 81行生成了一个很小的正态分布随机噪声,这个噪声是用来产生83行梯度计算的扰动的。83行计算了SynthesisNetwork生成的图像对的MappingNetwork生成的编码的梯度,并乘以了
pl_noise
作为随机扰动。由于create_graph
设为了True,所以这个算出来的梯度项也是可以用来计算损失并backward的,会根据二阶导来更新SynthesisNetwork的参数。这里得到的pl_grads
的shape和gen_ws
的shape是一样的,都是[batch_size, num_ws, dim_w]
- 84计算了
pl_grads
每个sample对不同w的平均向量二范数,得到的shape是[batch_size, ]
- 85行利用求得的
pl_lengths
对pl_mean
进行移动平均,pl_decay
是0.01,即是pl_mean = pl_mean + 0.01 * (pl_lengths - pl_mean)
,这里pl_mean
初始值是0,所以随着训练的迭代,pl_mean
会是一个保存了历次迭代的pl_lengths
的移动平均。 - 86行把更新后的
pl_mean
从梯度图中分离出来,以防止被梯度更新改变值,这个变量只是用来保存pl_lengths
的移动平均的,不应该被其它过程更新参数。 - 87行计算了一个
pl_penalty
变量,这个变量就是Greg阶段的损失了。所以可以看出,Greg阶段主要是惩罚pl_grads
的变化。89行pl_weight
是2,92行gain在这一阶段是4 - 95-104行Dmain阶段和Gmain阶段的区别仅在于,logit的符号不再取反,此时训练的是
Discriminator,需要它输出正确的结果,gain在这一阶段是1 - 109-131行在Dmain阶段和Dreg阶段都会运行。109行写了个很绕的表达式,其实就是,Dmain阶段
name
为'Dreal'
,Dreg阶段name
为'Dr1'
- 110-114行为Discriminator送入了真实图片,所以118-119行(仅在Dmain阶段运行)计算了对真实图片的discriminator loss,加负号是表示需要Discriminator预测的值越高越好,因为高代表real,低代表fake。
- 123-128行仅在Dreg阶段运行。124行同样计算了Discriminator的输出对输入的真实图片的梯度,但这里则直接以梯度值的平方和作为损失
- 131行gain是16
- 损失函数到此介绍完,值得注意的是Greg阶段和Dreg阶段都取了网络的输出对输入的梯度来计算损失,这看起来与机器学习中以参数W的范数作为损失项的一种正则化方法有点类似,因为网络的输出对输入的梯度可以大体可以视作是网络的参数。
torch_utils/ops
torch_utils/ops/persistence.py/persistent_class
- 这个函数在torch_utils/persistence.py文件的35行被定义。可以看到99行和130行,这个函数返回的是输入类的一个子类,这个子类为这个输入的类添加了一些功能,包括:保存类的初始化参数,为类添加打包函数(__reduce__方法的功能即为当代码被pickle打包时能输出正确的字符串等,展开说有点离题,具体可以自己去查,这里是暂时不需要了解的细节)
如何对这类代码进行修改以实现自己的idea
- 这类编程范式常见于微软和谷歌等大厂的完备API库,如mmdet等,他们提供了许多API接口和现有模型,也有API文档,代码层层包装,在接口外对代码进行修改是很困难的事情。但这些代码一般都提供了方便的自定义接口,只要找对方式,在这些库上实现自己的idea是很轻松的事情。
- 虽然个人对编程范式研究不是很多,但代码看得比较多,这类库个人感觉目的就是最大程度地避免对已有的函数和类的修改。也就是说,你如果想要在这上面实现自己的idea,你应该通过新建的方式而非修改的方式。无论你是想要实现新模型、新训练过程、新损失函数、新数据集、新增强手段等,在深度学习的全环节,这些过程都被包装成一个个的类,而在主函数中通过字符串来索引这些类。因此想要对某个环节进行替换,只需要两个步骤:
- 在对应的文件里新增该类的定义和实现。如你想把模型换成自己的模型,那就在network.py里面声明和定义自己的类。
- 在config或者命令行参数中,通过字符串形式指定某环节使用的类的名字。比如你定义了一个叫mynewmodel的类,那么你可能是通过将命令行调用改写为python train.py --model mynewmodel这样的方式实现对自己模型的训练。亦或者,有些代码将参数全部从jsonyaml文件读取,那么你就应该新建自己的json/yaml文件,在其中将模型名改为自己的模型,然后通过python train.py --config myconfig.yaml等方式来运行代码。
- 总而言之,对这类代码的常规修改方式就应该是通过添加自己的类,然后修改config指向自己的类,来实现。当然自己的类的定义一般不是随心所欲的,通常需要遵循某些规则,比如要定义某些方法,或是在某些方法内提供特定格式的返回值,等等。
- 通常来说,常规的idea都能通过代码提供的这些常规接口实现,当然也会出现常规接口无法实现的时候,一般就是你创新性地需要利用到某些信息,而这份库恰好没有把这些信息从包装里面传出来,这时候无可避免就得对代码进行修改而非新建。不过通常这些情况非常少,所以建议能不修改代码还是不修改代码,否则会带来很多意想不到的bug和麻烦。
TODO:
torch_utils/ops/conv2d_resample.py/conv2d_resample
torch_utils/ops/bias_act.py/bias_act
training.augment.AugmentPipe
metric_main
misc.assert_shape
misc.profiled_function
upfirdn2d.upsample2d
upfirdn2d.downsample2d
文章来源:https://www.toymoban.com/news/detail-409469.html
到了这里,关于StyleGAN2代码阅读笔记的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!