轻轻松松使用StyleGAN2(七):详解并中文注释training_loop.py,看一看StyleGAN2是怎样构建训练图并进行训练的?

StyleGAN2本质上是通过假图像生成器generator与真图像判别器discrimnator之间的对抗,最终使判别器无法判别真假(对于由假图像和真图像共同组成的数据集,判别器给出正确标签的概率为 50%)。其过程表现为两个神经网络的权重和偏置不断调整,使得对于生成器生成的假图像,判别器判别为假的概率最小,即:生成器神经网络的运算矩阵所表达的特征期望平均值(度量标准包括:FID、PPL、LPIPS等)逼近真实图像样本的平均值,且特征期望方差为最小;同时判别器对真假图像的混合数据集能给出正确标签(即:判定真图像为真,假图像为假)的概率最大。

事实上,通过阅读StyleGAN2的源代码,我们可以看到这个过程大致是这样的:

(1)随机生成一批(minibatch_size)潜码 latents,通过生成器 G 得到一批假图像 fake_images_out;

(2)再从训练数据集中取一部分(minibatch)真实图像 reals;

(3)将真、假图像分别交给判别器 D 去计算得分(即:判定为真的概率,分别为 real_scores_out 和 fake_scores_out),并计算这些得分的全局交叉熵( -log(1-sigmoid(fake_scores_out)) -log(sigmoid(real_scores_out)) )——假图像判别为假的概率的对数与真图像判别为真的概率的对数之和的相反数,把全局交叉熵与正则化项求和作为判别器 D 的损失函数;

(4)同时,生成器 G 的损失函数则只考虑它自身的交叉熵( -log(sigmoid(fake_scores_out)) )——假图像判别为真的概率的对数的相反数,再与PPL(感知路径长度)正则化项求和作为生成器G的损失函数;

(5)优化的过程就是通过梯度下降,使得 D 和 G 的损失函数取得最小值。

生成对抗网络(GAN)的发明人 Goodfellow 称 G 采用这样的损失函数的博弈为 Non-saturating heuristic game ——在实践中这样做很有效,但并不保证 GAN 一定会收敛达到纳什均衡。

这里的所谓“对抗”,是说 G 的损失函数要使其生成的假图像被 D 判别为真的概率最大(即:判别为假的概率最小),而与此同时,D 的损失函数的一部分要求假图像被判别为假的概率最大,G 和 D 同时向相反方向优化,但有可能找到唯一解。

StyleGAN2训练神经网络的主程序是.\run_training.py,这个程序里,大量的内容都是针对config-a、config-b、config-c ......config-e进行参数配置,注释中说:

# Configs A-E: Shrink networks to match original StyleGAN.

大致意思是说,从A-E是在baseline基础版本上不断添加优化的“缩水版”的StyleGAN2网络。

本文研究的是F,即最终确定的效果最佳的“大网络”。从A-F逐步优化的过程,代码中的注释如下:

    'config-a', # Baseline StyleGAN
    'config-b', # + Weight demodulation
    'config-c', # + Lazy regularization
    'config-d', # + Path length regularization
    'config-e', # + No growing, new G & D arch.
    'config-f', # + Large networks (default)

即:

配置A:StyleGAN基线模型

配置B:+ 权重解调(用于消除液滴伪影)

配置C:+ 延迟正则化(用于降低正则化计算开销)

配置D:+ 路径长度正则化(用于提升图像质量)

配置E:+ 无(渐进式)生长,新生成器/判别器架构(用于消除阶段性伪影)

配置F:+ 大网络(缺省)(用于改进1024x1024分辨率的图像质量)

.\run_training.py的代码内容很单薄,正如春天的溪水一样清浅,我们这里不做过多解读,我们下面重点看一下.\run_training.py调用的核心训练函数代码.\training\training_loop.py,简单研究一下StyleGAN2是怎样构建网络并训练数据的。

我们先总结几条内容,如下:

(一)StyleGAN2使用“TFRecordDataset类”装载训练数据,TFRecordDataset类用于从.tfrecords文件中加载不同lod(levels of detail,以1024x1024的图像为例,从0--8)的shape、labels(标签)和图像数据,并且定义了这些数据的预处理方法,相应的内容我们在上一篇文章里已经讲过,内容请参考:

轻轻松松使用StyleGAN2(六):StyleGAN2 Encoder是怎样加载训练数据的?源代码+中文注释,dataset_tool.py和dataset.py

(二)StyleGAN2使用.dnnlib.tflib作为构建和训练神经网络的基础架构,包括定义网络、设定优化器、设定损失函数、定义(注册)梯度、运行迭代操作等——dnnlib.tflib是英伟达团队在tensorflow上构建的一组库函数,以便于代码维护,也有助于理解代码结构;

(三)StyleGAN2的核心网络结构在.\training\networks_stylegan2.py中定义,包括:多个版本的生成器、判别器等,生成器和判别器是两个独立的神经网络,生成器基于随机向量生成假图像(训练使得生成器捕捉到真实图像的特征,以尽可能生成令判别器判定为真的假图像),判别器用于判别图像是否为真(即:是不是来自于训练数据集中的真实图像),各自进行训练优化;

(四)StyleGAN2的损失函数在.\training\loss.py中定义,生成对抗网络中的“对抗”,主要体现在生成器和判别器各自损失函数的定义上,由此造成两个神经网络迭代优化时的对抗与竞合。

下面是.\training\training_loop.py的流程图(主要节点一共分为10步):

.\training\training_loop.py源代码的中文注释如下:

# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html

"""Main training script."""

import numpy as np
import tensorflow as tf
import dnnlib
import dnnlib.tflib as tflib
from dnnlib.tflib.autosummary import autosummary

from training import dataset
from training import misc
from metrics import metric_base

#----------------------------------------------------------------------------
# Just-in-time processing of training images before feeding them to the networks.
# 在将图像数据投喂给(神经)网络前,对用于训练的图像即时进行(预)处理

def process_reals(x, labels, lod, mirror_augment, drange_data, drange_net):
    # 将向量数据的取值范围从(0,255)调整到(-1,1)
    with tf.name_scope('DynamicRange'):
        x = tf.cast(x, tf.float32)
        x = misc.adjust_dynamic_range(x, drange_data, drange_net)
    # (临时地)生成均匀分布的随机向量,shape为[tf.shape(x)[0]],数值默认为[0,1]之间
    # 如果这个临时向量的元素值小于0.5,则x的向量取值不变;否则对应位置的x向量的元素值在axis=[3]这一层颠倒顺序
    # tf.shape(x)应为[minibatch, channels, height, weight]
    if mirror_augment:
        with tf.name_scope('MirrorAugment'):
            x = tf.where(tf.random_uniform([tf.shape(x)[0]]) < 0.5, x, tf.reverse(x, [3]))
    # Smooth crossfade between consecutive levels-of-detail.
    # 在连续的lod之间实现平滑的交叉渐入渐出
    with tf.name_scope('FadeLOD'):
        s = tf.shape(x)
        # 高、宽减半,reshape()
        # 举例,对于1024x1024的图片,相当于分成4片,每片512x512
        y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
        # 减半的图像在第3维、第5维两个维度求平均值,保持6个维度
        # 即:对于4片数据上对应像点的4个数据取平均值,保存为512x512的lod数据
        y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
        # 扩展数据,将第3维和第5维数据加倍,恢复为4片512x512的数据
        y = tf.tile(y, [1, 1, 1, 2, 1, 2])
        # reshape,恢复到tf.shape(x),但相邻4个像点的数据用它们的平均值代替
        y = tf.reshape(y, [-1, s[1], s[2], s[3]])
        # 在x、y两个向量之间进行线性插值运算
        # tf.floor(lod):取比lod小的最大整数
        x = tflib.lerp(x, y, lod - tf.floor(lod))
    # Upscale to match the expected input/output size of the networks
    # 将x向量代表的图像的高和宽各扩大为原来的factor倍
    # factor = 2 ** tf.floor(lod),从1024x1024到4x4,对应的lod分别为0,1,2,3,4,5,6,7,8
    with tf.name_scope('UpscaleLOD'):
        s = tf.shape(x)
        factor = tf.cast(2 ** tf.floor(lod), tf.int32)
        x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
        x = tf.tile(x, [1, 1, 1, factor, 1, factor])
        x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
    return x, labels

#----------------------------------------------------------------------------
# Evaluate time-varying training p
评论 12
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值