StyleGAN2代码阅读笔记-创新互联

源代码地址：https://github.com/NVlabs/stylegan2-ada-pytorch

创新互联建站始终坚持【策划先行，效果至上】的经营理念，通过多达十年累计超上千家客户的网站建设总结了一套系统有效的全网推广解决方案，现已广泛运用于各行各业的客户，其中包括：石雕等企业，备受客户好评。

• 这是一篇代码阅读笔记，顾名思义是对代码进行阅读，讲解的笔记。对象是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_modconvbool变量，后续用到的时候再展开，这里只需要知道只有测试的时候才有可能是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，然后stylesreshape成[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]的向量，会自动广播，其实就是把stylesrepeat成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行把xreshape成[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

你是否还在寻找稳定的海外服务器提供商？创新互联www.cdcxhl.cn海外机房具备T级流量清洗系统配攻击溯源，准确流量调度确保服务器高可用性，企业级服务器适合批量采购，新人活动首月15元起，快前往官网查看详情吧

标题名称：StyleGAN2代码阅读笔记-创新互联
浏览路径：http://cdiso.cn/article/dggjed.html