0
点赞
收藏
分享

微信扫一扫

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战


大家好,欢迎来到专栏《百战GAN》,在这个专栏里,我们会进行GAN相关项目的核心思想讲解,代码的详解,模型的训练和测试等内容。

作者&编辑 | 言有三

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_计算机视觉

本文资源与生成结果展示

本文篇幅:9000字

背景要求:会使用Python和Pytorch

附带资料:参考论文和项目

1 项目背景

GAN自从被提出来后,技术发展就非常迅猛,已经被落地于众多的方向,其应用涉及图像与视频生成,数据仿真与增强,各种各样的图像风格化任务,人脸与人体图像编辑,图像质量提升。

其中GAN最早期也是最经典的任务,就是高质量图像生成,当前已经可以生成1024分辨率以上的高清逼真图像,如下图生成了一些假明星脸。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_计算机视觉_02

以上图片就是使用StyleGAN进行生成,StyleGAN系列模型是当前最优秀的图像生成框架,不仅被用于图像生成领域,也被用于其他诸如图像修复等方向,是生成对抗网络必须掌握的内容。

2 原理简介

StyleGAN[1]是一个强大的可以控制生成图片属性的框架,它采用了全新的生成模型,分层的属性控制,Progressive GAN的渐进式分辨率提升策略,能够生成1024×1024分辨率的人脸图像,并且可以进行属性的精确控制与编辑, 下图展示了StyleGAN论文中生成的人脸图片。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_算法_03

StyleGAN与传统的生成器的对比结构如下图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_人工智能_04

接下来我们对StyleGAN的原理进行解读。

2. 映射网络f

映射网络f总共有8层全连接层,输入是512维的噪声向量Z,经过8个全连接层,得到512维的潜在空间向量W,这样编码的好处是为了摆脱输入向量受输入数据集分布的影响,下面参考论文中的简单案例进行说明,如下图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_算法_05

训练数据集通常是有偏的,比如在人脸的属性中,性别包括男女,头发包括长短,其中{男,长发}属性一起出现的概率较低,而{男,短发},{女,长发},{女,短发}一起出现的概率较高,反映到空间中就是一个不均匀的分布,如图(a)。

如果我们仅仅使用随机采样的噪声向量Z来映射,因为噪声Z的分布在全空间,为了拟合训练数据集,必定存在不均匀的映射区域,如图(b),这增加了从Z到生成图片的模型学习难度,因为属性之间的耦合关系非常复杂。

假如通过映射网络f首先对Z进行映射得到W,不仅可以保证与训练集一致的分布,还获得更加均匀的属性分布,潜在向量空间W与生成图片的属性之间有更好的线性关系,这有利于对生成图片的属性控制,因此W更加合适作为生成器的输入。

2.2 生成网络g

接下来我们再看生成网络g,它通过分层的控制来实现不同粒度人脸属性的编辑。

AdaIN层是一个在生成对抗网络和风格化领域中应用非常广泛的归一化层,在风格编码任务中,它可以替换批归一化层(BN)获得更好的结果,其定义如下;

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_机器学习_06

AdaIN 的具体实现过程是:将512维的向量W通过一个可学习的仿射变换,生成缩放因子与偏差因子,这两个因子会与实例标准化(即Instance Normalization,简称IN)之后的输出做加权求和,原理示意如下图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_python_07

后来StyleGAN的研究者发现,对不同的AdaIN层使用不同的W向量是有益的,因此W的维度被拓展成18×512,称之为W',其中18对应AdaIN层的数量。

由于实例标准化对每个特征图单独计算,尺度和偏移的维度也与特征图通道数有关。通过缩放因子与偏差因子,我们可以实现图片的整体样式控制,所以它们可以被称之为风格向量。

生成网络synthesis network g是一个分辨率逐级提升的结构,总共有17个卷积层,除了第1层以外,每两层上采样一个尺度,分辨率从4×4提升到1024×1024,训练方式与Progressive GAN相同。每一级分辨率都有两个AdaIN层,我们可以将其称为1个风格化模块,一共9个风格化模块。

以StyleGAN生成的人脸图像为例,作者在论文的实验中发现,按照尺度可以将人脸特征分为3个层级,全局特征,中级特征与细节特征,如下图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_计算机视觉_08

全局特征由分辨率不超过8×8的风格化模块控制,主要包括面部姿势、发型、面部形状等特征。

中级特征由分辨率在16×16和32×32的风格化模块控制,主要包括更精细的面部特征、发型、眼睛的睁闭等。

细节特征由分辨率从64×64到1024×1024的风格化模块控制,主要包括眼睛、头发和皮肤等纹理和颜色细节。

另外在每1个风格化模块的卷积层之后,AdaIN层之前,都添加了通道特征图级别的高斯噪声,每一层各个通道的噪声输入共用,但是需要乘以可学习的权重后再添加到特征图中。噪声的添加可以对更加细微的生成结果进行随机控制,增强生成图片的模式丰富性,相关实验结果可以看下面的实践。

因为StyleGAN 生成图像的特征是由权重W和AdaIN层控制,所以生成器的初始输入不再需要输入噪声,而是用全1的常量值替代。

2.3 训练技巧

StyleGAN是一个非常优秀的生成架构,但仅仅依靠优良的架构并不足以取得非常高质量的生成结果,还需要一些训练技巧辅助模型的训练,主要包含两个,样式正则化(即mixing regularization)与W向量截断。

为了降低StyleGAN生成器中各个级别特征的相关性,StyleGAN采用了样式正则化(mixing regularization)训练技巧。它通过在训练的时候,随机选择两个输入向量Z1和Z2,经过映射网络得到中间向量W1,W2,然后随机交换W1和W2的部分内容,从而实现两幅图像风格的交换。

如下图中向量a分为a1和a2两段,向量b分别b1和b2两段,将a1和b2组合成一个新的与a和b长度相同的向量,就是一种常用的样式向量混合。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_算法_09

在论文中作者发现,将B域的风格在4~8粗粒度(小尺度)往A域进行迁移时,结果会保留B域的发型,脸型等全局信息,而颜色以及纹理来来自于A域。

将B域的风格在16~32中等粒度往A域进行迁移时,结果会保留B域小尺度的脸部细节,如发型,眼神等,而姿态等全局信息则来自于A域。

将B域的风格在64~1024细粒度(大尺度)往A域进行迁移时,结果会保留B域的一些细节纹理和颜色风格,其他都来自A域。

另外一个重要的技巧就是W向量的截断技巧,具体的做法是首先对W向量计算出统计均值,然后通过截断函数来生成新的W向量,如下式:

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_计算机视觉_10

其中截断函数的值域是(-1,1)。

2.4 StyleGAN的评估

StyleGAN额外提出了两个新的评估方法,包括感知路径长度perceptual path length和线性可分性Linear separability。

路径长度评估的是潜在空间Z或者W中端点的平均距离,具体计算为训练过程中相邻时间节点上的两个生成图像的距离,基于Z的定义如下式,基于W的定义方法类似:

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_python_11

其中Slerp表示spherical interpolation,是一种球形空间的采样方法;d表示VGG特征空间的L1距离;t表示某一个时间点,表示相邻的时间步。

一个非常好的潜在空间向量,在空间中应该是线性分布的,即沿着某一个路径,可以编辑相关属性,当我们想要生成特定属性的图像时,在该路径上进行采样是最高效的,比如下图中的绿色虚线路径,可以在任意节点采样生成‘猫’图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_人工智能_12

而蓝色虚线表示一条更长的路径,虽然在路径的终点进行采样可以生成满足属性的图片,但是其中间节点采样得到的向量不能够生成‘猫’,因此蓝色虚线路径质量不如绿色虚线。它们的质量差异在图中的直观表达就是路径的长度,即perceptual path length,更短的路径表示质量更高的空间映射。

另一个评估指标即线性可分性Linear separability,它用于评估Latent向量是否具有足够的属性可分类性。

首先我们使用分布z ∼ P(z)生成200000张图片,然后对其训练1个CNN图片分类器得到某一个属性的二分类标签,比如是否微笑;

接下来我们对潜在空间向量Z或者W使用SVM进行分类,计算条件熵H(Y|X),其中X是SVM分类器结果,Y是CNN图片分类器结果。

下图左展示了微笑与不微笑两类样本,右图展示了属性向量的空间分布。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_算法_13

假如潜在空间中的属性向量具有良好的线性可分性,H(Y|X)就会越小,它表示需要多少额外信息来决定类别,更低的值反映出更可分的属性向量分布。

3 模型解读

接下来我们来进行代码实战,使用预训练模型进行测试,对StyleGAN核心的模型代码进行解读。

3.1 生成器定义

首先我们来看图像生成网络(synthesis network)的定义:

## synthesis network定义

class Generator(nn.Module):

    def __init__(self, code_dim, fused=True):

        super().__init__()

        ##  9个尺度的卷积block,从4×4到64×64,使用双线性上采样;从64×64到1024×1024,使用转置卷积进行上采样

        self.progression = nn.ModuleList(

            [

                StyledConvBlock(512, 512, 3, 1, initial=True),  ##4×4

                StyledConvBlock(512, 512, 3, 1, upsample=True),  ## 8×8

                StyledConvBlock(512, 512, 3, 1, upsample=True),  ## 16×16

                StyledConvBlock(512, 512, 3, 1, upsample=True),  ## 32×32

                StyledConvBlock(512, 256, 3, 1, upsample=True),  ## 64×64

                StyledConvBlock(256, 128, 3, 1, upsample=True, fused=fused),  ## 128×128

                StyledConvBlock(128, 64, 3, 1, upsample=True, fused=fused),  ## 256×256

                StyledConvBlock(64, 32, 3, 1, upsample=True, fused=fused),  ## 512×512

                StyledConvBlock(32, 16, 3, 1, upsample=True, fused=fused),  ## 1024×1024

            ]

        )

        ## 9个尺度的1×1 to_rgb卷积层,将特征图输出为RGB图片,与9个风格模块对应

        self.to_rgb = nn.ModuleList(

            [

                EqualConv2d(512, 3, 1),

                EqualConv2d(512, 3, 1),

                EqualConv2d(512, 3, 1),

                EqualConv2d(512, 3, 1),

                EqualConv2d(256, 3, 1),

                EqualConv2d(128, 3, 1),

                EqualConv2d(64, 3, 1),

                EqualConv2d(32, 3, 1),

                EqualConv2d(16, 3, 1),

            ]

        )

    def forward(self, style, noise, step=0, alpha=1, mixing_range=(-1, 1)):

        out = noise[0] ## 取噪声向量为输入

        if len(style) < 2: ## 输入只有1个风格向量,表示不进行样式混合,inject_index=10

            inject_index = [len(self.progression) + 1]

        else:

            ## 不止一个style向量,可以进行样式混合训练,生成长度为len(style) - 1))的样式混合交叉点序列,其数值大小不超过step

            inject_index = sorted(random.sample(list(range(step)), len(style) - 1))

        crossover = 0 ##用于样式混合的位置

        for i, (conv, to_rgb) in enumerate(zip(self.progression, self.to_rgb)):

            if mixing_range == (-1, 1):

## 根据前面生成的随机数,来决定样式混合的index

                if crossover < len(inject_index) and i > inject_index[crossover]:

                    crossover = min(crossover + 1, len(style))

                style_step = style[crossover] ##获得交叉的style起始点

            else:

## ## 根据mixing_range来觉得样式混合的区间,mixing_range[0] <= i <= mixing_range[1]取style[1],其他取style[0]

                if mixing_range[0] <= i <= mixing_range[1]:

                    style_step = style[1] ## 取第2个样本样式

                else:

                    style_step = style[0] ## 取第1个样本样式

            if i > 0 and step > 0:

                out_prev = out

            ## 将噪声与风格向量输入风格模块

            out = conv(out, style_step, noise[i])

            if i == step: ## 最后1级分辨率,输出图片

                out = to_rgb(out) ## 1×1卷积

                ## 最后结果是否进行alpha融合

                if i > 0 and 0 <= alpha < 1:

                    skip_rgb = self.to_rgb[i - 1](out_prev) ##获得上一级分辨率结果进行2倍上采样

                    skip_rgb = F.interpolate(skip_rgb, scale_factor=2, mode='nearest')

                    out = (1 - alpha) * skip_rgb + alpha * out

                break

        return out

首先可以看到总共包含了9个风格模块,即StyledConvBlock,其中第1个风格模块不需要进行上采样,剩下8个模块需要进行上采样。每1个风格模块都对应1个to_rgb卷积层,可以输出当前分辨率的图像。

风格模块的输入包括了噪声向量和风格向量,接下来我们解读风格模块:

从中可以看到,除了第1个风格层输出4×4×512大小的值为全1的常量特征图,其他都需要进行上采样,对于128及以上的分辨率使用转置卷积上采样,对于128以下的分辨率使用最近邻上采样。

其中ConstantInput定义如下:

class ConstantInput(nn.Module):

    def __init__(self, channel, size=4):

        super().__init__()

        self.input = nn.Parameter(torch.randn(1, channel, size, size))

    def forward(self, input):

        batch = input.shape[0]

        out = self.input.repeat(batch, 1, 1, 1)

        return out

转置卷积上采样定义如下:

## 转置卷积上采样,其中权重参数自己定义

class FusedUpsample(nn.Module):

    def __init__(self, in_channel, out_channel, kernel_size, padding=0):

        super().__init__()

        weight = torch.randn(in_channel, out_channel, kernel_size, kernel_size)

        bias = torch.zeros(out_channel)

        fan_in = in_channel * kernel_size * kernel_size ##神经元数量

        self.multiplier = sqrt(2 / fan_in)

        self.weight = nn.Parameter(weight)

        self.bias = nn.Parameter(bias)

        self.pad = padding

    def forward(self, input):

        weight = F.pad(self.weight * self.multiplier, [1, 1, 1, 1])

        weight = (

            weight[:, :, 1:, 1:]

            + weight[:, :, :-1, 1:]

            + weight[:, :, 1:, :-1]

            + weight[:, :, :-1, :-1]

        ) / 4

        out = F.conv_transpose2d(input, weight, self.bias, stride=2, padding=self.pad)

        return out

噪声模块的定义如下,它通过权重和图像进行相加融合:

## 添加噪声,噪声权重可以学习

class NoiseInjection(nn.Module):

    def __init__(self, channel):

        super().__init__()

        self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))

    def forward(self, image, noise):

        return image + self.weight * noise

AdaIN模块的定义如下,它通过缩放和偏置系数控制风格:

## 自适应的IN层

class AdaptiveInstanceNorm(nn.Module):

    def __init__(self, in_channel, style_dim):

        super().__init__()

        self.norm = nn.InstanceNorm2d(in_channel) ##创建IN层

        self.style = EqualLinear(style_dim, in_channel * 2) ##全连接层,将W向量变成AdaIN层系数S

        self.style.linear.bias.data[:in_channel] = 1

        self.style.linear.bias.data[in_channel:] = 0

def forward(self, input, style):

## 输入style为风格向量W,长度为512;经过self.style得到输出风格矩阵S,通道数等于输入通道数的2倍

        style = self.style(style).unsqueeze(2).unsqueeze(3)

        gamma, beta = style.chunk(2, 1) ## 获得缩放和偏置系数,按1轴(通道)分为2部分

        out = self.norm(input) ##IN归一化

        out = gamma * out + beta

        return out

Style向量需要通过仿射变换从W向量中学习,EqualLinear定义如下:

## 全连接层

class EqualLinear(nn.Module):

    def __init__(self, in_dim, out_dim):

        super().__init__()

        linear = nn.Linear(in_dim, out_dim)

        linear.weight.data.normal_()

        linear.bias.data.zero_()

        self.linear = equal_lr(linear)

    def forward(self, input):

        return self.linear(input)

EqualLinear层的输入维度是style_dim,即512,输出是in_channel * 2,其中乘以2是因为缩放和偏置系数要产生两份,而in_channel对应的就是要作用的通道的数量。

在上述代码中我们可以看到不管是卷积层还是全连接层,都需要调用equal_lr函数进行权重的归一化,这是StyleGAN的训练工程技巧之一,它根据当前层的神经元数量,对权重进行归一化,从而实现让各层有等价学习率的效果,equal_lr函数的实现如下。

## 归一化学习率

class EqualLR:

    def __init__(self, name):

        self.name = name

    def compute_weight(self, module):

        weight = getattr(module, self.name + '_orig')

        ## 输入神经元数目,每一层卷积核数量=Nin*Nout*K*K,

        fan_in = weight.data.size(1) * weight.data[0][0].numel()

        return weight * sqrt(2 / fan_in)

    @staticmethod

    def apply(module, name):

        fn = EqualLR(name)

        weight = getattr(module, name)

        del module._parameters[name]

        module.register_parameter(name + '_orig', nn.Parameter(weight.data))

        module.register_forward_pre_hook(fn)

        return fn

    def __call__(self, module, input):

        weight = self.compute_weight(module)

        setattr(module, self.name, weight)

def equal_lr(module, name='weight'):

    EqualLR.apply(module, name)

return module

完整的生成器定义如下:

## 完整的生成器定义

class StyledGenerator(nn.Module):

    def __init__(self, code_dim=512, n_mlp=8):

        super().__init__()

        self.generator = Generator(code_dim) ## synthesis network

        ## mapping network定义,包含8个全连接层,n_mlp=8

        layers = [PixelNorm()]

        for i in range(n_mlp):

            layers.append(EqualLinear(code_dim, code_dim))

            layers.append(nn.LeakyReLU(0.2))

        ## mapping network f,用于从噪声向量Z生成Latent向量W(即风格向量)

        self.style = nn.Sequential(*layers)

    def forward(

        self,

        input, ##输入向量Z

        noise=None, ##噪声向量,可选的

        step=0, ##上采样因子

        alpha=1, ##融合因子

        mean_style=None, ##平均风格向量W

        style_weight=0, ##风格向量权重

        mixing_range=(-1, -1), ##混合区间变量

    ):

        styles = [] ##风格向量W

        if type(input) not in (list, tuple):

            input = [input]

        for i in input:

            styles.append(self.style(i)) ## 调用mapping network,生成第i个风格向量W

        batch = input[0].shape[0] ## batchsize大小

        if noise is None:

            noise = []

            for i in range(step + 1): ## 0~8,共9层noise

                size = 4 * 2 ** i ## 每一层的尺度,第一层为4*4,每一层的各个通道共用噪声

                noise.append(torch.randn(batch, 1, size, size, device=input[0].device))

## 基于平均风格向量和当前生成的风格向量,获得完整的风格向量

        if mean_style is not None:

            styles_norm = [] ## 风格数组[1*512]

            for style in styles:

                styles_norm.append(mean_style + style_weight * (style - mean_style))

            styles = styles_norm

        return self.generator(styles, noise, step, alpha, mixing_range=mixing_range)

以上就是生成器的主要代码,接下来我们再看判别器的定义。

3.2 判别器定义

判别器也采用了Progressive GAN中渐进式的判别结构,定义如下:

class Discriminator(nn.Module):

    def __init__(self, fused=True, from_rgb_activate=False):

        super().__init__()

        self.progression = nn.ModuleList(

            [

                ConvBlock(16, 32, 3, 1, downsample=True, fused=fused),  ## 512×512

                ConvBlock(32, 64, 3, 1, downsample=True, fused=fused),  ## 256×256

                ConvBlock(64, 128, 3, 1, downsample=True, fused=fused),  ## 128×128

                ConvBlock(128, 256, 3, 1, downsample=True, fused=fused),  ## 64×64

                ConvBlock(256, 512, 3, 1, downsample=True),  ## 32×32

                ConvBlock(512, 512, 3, 1, downsample=True),  ## 16×16

                ConvBlock(512, 512, 3, 1, downsample=True),  ## 8×8

                ConvBlock(512, 512, 3, 1, downsample=True),  ## 4×4

                ConvBlock(512, 512, 3, 1, 4, 0),

            ]

        )

        ## 从RGB图片转为概率

        def make_from_rgb(out_channel):

            if from_rgb_activate:

                return nn.Sequential(EqualConv2d(3, out_channel, 1), nn.LeakyReLU(0.2))

            else:

                return EqualConv2d(3, out_channel, 1)

        self.from_rgb = nn.ModuleList(

            [

                make_from_rgb(16),

                make_from_rgb(32),

                make_from_rgb(64),

                make_from_rgb(128),

                make_from_rgb(256),

                make_from_rgb(512),

                make_from_rgb(512),

                make_from_rgb(512),

                make_from_rgb(512),

            ]

        )

        self.n_layer = len(self.progression)

        self.linear = EqualLinear(512, 1)

    def forward(self, input, step=0, alpha=1):

        for i in range(step, -1, -1):

            index = self.n_layer - i - 1

            if i == step: ##最高级,输入图片

                out = self.from_rgb[index](input)

            if i == 0:

                out_std = torch.sqrt(out.var(0, unbiased=False) + 1e-8)

                mean_std = out_std.mean()

                mean_std = mean_std.expand(out.size(0), 1, 4, 4)

                out = torch.cat([out, mean_std], 1)

            out = self.progression[index](out)

            ## 判别器的相邻层融合

            if i > 0:

                if i == step and 0 <= alpha < 1:

                    skip_rgb = F.avg_pool2d(input, 2)

                    skip_rgb = self.from_rgb[index + 1](skip_rgb)

                    out = (1 - alpha) * skip_rgb + alpha * out

        out = out.squeeze(2).squeeze(2)

        out = self.linear(out)

        return out

首先可以看到总共包含了9个卷积模块,即ConvBlock,其中第9个风格模块不需要进行下采样,剩下8个模块需要进行下采样。每1个风格模块都对应1个make_from_rgb卷积层,可以根据当前分辨率的图像输出真假预测概率。

ConvBlock模块的定义如下:

class ConvBlock(nn.Module):

    def __init__(

        self,

        in_channel,

        out_channel,

        kernel_size,

        padding,

        kernel_size2=None,

        padding2=None,

        downsample=False,

        fused=False,

    ):

        super().__init__()

        pad1 = padding

        pad2 = padding

        if padding2 is not None:

            pad2 = padding2

        kernel1 = kernel_size

        kernel2 = kernel_size

## 最后一层kernel_size2=4,其他层输入为none

        if kernel_size2 is not None:

            kernel2 = kernel_size2

        self.conv1 = nn.Sequential(

            EqualConv2d(in_channel, out_channel, kernel1, padding=pad1),

            nn.LeakyReLU(0.2),

        )

        if downsample:

            if fused: ## 对于128及以上的分辨率,使用步长为2的卷积

                self.conv2 = nn.Sequential(

                    Blur(out_channel),

                    FusedDownsample(out_channel, out_channel, kernel2, padding=pad2),

                    nn.LeakyReLU(0.2),

                )

            else: ## 对于64及以下的分辨率,使用平均池化

                self.conv2 = nn.Sequential(

                    Blur(out_channel),

                    EqualConv2d(out_channel, out_channel, kernel2, padding=pad2),

                    nn.AvgPool2d(2),

                    nn.LeakyReLU(0.2),

                )

        else:

            self.conv2 = nn.Sequential(

                EqualConv2d(out_channel, out_channel, kernel2, padding=pad2),

                nn.LeakyReLU(0.2),

            )

    def forward(self, input):

        out = self.conv1(input)

        out = self.conv2(out)

        return out

与生成器中对不同分辨率模块采用不同上采样方法的策略类似,对于128及以上的分辨率使用带步长的卷积进行下采样,对于128以下的分辨率使用平均池化进行下采样,具体的代码请读者阅读完整工程。

4 图片生成实验

接下来我们进行人脸图像生成,首先我们需要根据开源项目中的提示下载相关的预训练模型,本次我们下载1024分辨率的生成模型,然后使用该预训练模型来生成图像。

4.1 人脸生成

首先我们构建预测器并生成人脸,核心的推理代码如下:

## 构建预测器

class Predictor():

    def __init__ (self,modelpath):

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.generator = StyledGenerator(512).to(self.device)  ## 定义生成器

        ## 载入训练好的模型权重

        weights = torch.load(modelpath,map_location=self.device)       

        self.generator.load_state_dict(weights["generator"])

        self.generator.eval()  ## 设置推理模式

        ## 获得平均风格向量

        self.mean_style = get_mean_style(self.device)

     ## 预测函数

     def predict(self, seed, output_path):

        torch.manual_seed(seed)  ## 为CPU设置种子用于生成随机数,使得结果确定

        step = int(math.log(SIZE, 2)) – 2

nsamples = 15

        img = self.generator(

            torch.randn(nsamples, 512).to(self.device),

            step=step,

            alpha=1,

            mean_style=self.mean_style,

            style_weight=0.7,

        )

utils.save_image(img, output_path, normalize=True)

if __name__ == '__main__':

    modelpath = "checkpoints/stylegan-1024px-new.model"

predictor = Predictor(modelpath)

## 基于不同的随机种子,运行10次获得生成结果

    for i in range(0,10):

        predictor.predict(i,'results/'+str(i)+'.png')

在初始化函数init中定义了生成器,获得了平均风格向量,在predict函数中调用generator生成了图片。

其中平均风格向量的获取函数为:

## 平均风格向量获取

@torch.no_grad()

def get_mean_style(generator, device):

    mean_style = None

for i in range(100):

        ## 从随机向量Z,经过mapping network得到W

        style = generator.mean_style(torch.randn(1024, 512).mean(0, keepdim=True).to(device))

        if mean_style is None:

            mean_style = style

        else:

            mean_style += style

    mean_style /= 100

    return mean_style

核心代码为将1×512维的随机向量Z输入生成器generator中的mean_style函数,产生1×512维的向量W,总共统计100次的平均值,得到的结果就是mean_style向量。

generator每次生成n_sample个样本,输入参数包括随机向量Z,step,alpha,style_weight。

其中step是上采样次数因子,当生成图片的分辨率为1024时,它等于8。因为输入是4×4大小的图,要经过28=256倍上采样。

alpha是一个跳层连接的融合因子,用于融合不同层不同分辨率的特征,默认为1,表示不进行融合。

style_weight是截断权重,权重越大,生成的图片越偏离平均脸,权重为0,则会生成平均脸。

下图是当截断权重为0的生成结果,可以看出生成的主体没有发生变化,只有背景有微小变化,它来自于在synthesis network中添加的输入噪声的影响。假如我们想要生成完全一样的人脸,可以将随机种子固定。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_计算机视觉_14

下图是当截断权重为0.7的生成结果,可以看出生成的主体发生了变化,可以生成各种真实的人脸。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_机器学习_15

4.2 样式混合编辑

StyleGAN在训练的时候使用了样式混合来提供正则化,我们接下来查看样式混合的结果,核心代码如下。

## 样式混合

@torch.no_grad()

def style_mixing(generator, step, mean_style, n_source, n_target, device):

    ## 两个样式向量

    source_code = torch.randn(n_source, 512).to(device)

    target_code = torch.randn(n_target, 512).to(device)

    shape = 4 * 2 ** step ##1024分辨率

    alpha = 1

    images = [torch.ones(1, 3, shape, shape).to(device) * -1]

    ## 源域图

    source_image = generator(

        source_code, step=step, alpha=alpha, mean_style=mean_style, style_weight=0.7

    )

    ## 目标域图

    target_image = generator(

        target_code, step=step, alpha=alpha, mean_style=mean_style, style_weight=0.7

    )

    images.append(source_image) ##存储源域图

    ## 样式混合

    for i in range(n_target):

        image = generator(

            [target_code[i].unsqueeze(0).repeat(n_source, 1), source_code],

            step=step,

            alpha=alpha,

            mean_style=mean_style,

            style_weight=0.7,

            mixing_range=(0, 1),

        )

        images.append(target_image[i].unsqueeze(0)) ##存储目标域图

        images.append(image) ##存储混合样式图

    images = torch.cat(images, 0)

在上述代码中,首先根据n_source,n_target生成源域和目标域的图,然后逐个将各自的样式向量进行混合。

混合的方式由mixing_range决定,有两种混合方式。当mixing_range=(-1, 1)时,为随机混合方式,即随机选择两个向量的交换点。

当指定有效的范围时,则根据有效范围进行混合。对于1024分辨率的图片,总共有9个风格化层,对应4,8,16,32,64,128,256,512,1024共9级分辨率。因此有效的样式混合范围处于0~8之间,我们取(0,1)进行接下来的样式混合实验,根据代码可以知道它表示4,8分辨率的特征取自于第2个风格向量,16,32,64,128,256,512,1024分辨率的特征取自于第1个风格向量。

下图展示了样式混合的结果图。

【百战GAN】StyleGAN原理详解与人脸图像生成代码实战_人工智能_16

第1行表示源域图,第1列表示目标域图,其他表示源域图和目标域的样式混合结果图。可以看出样式混合图保留了源域图的姿态,发型,脸型等宏观属性,保留了目标域图中的肤色,眼睛,毛发纹理等微观特征,实现了逼真的样式混合。

本文参考的文献如下:

[1] Karras T, Laine S, Aila T. A style-based generator architecture for generative adversarial networks[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019: 4401-4410.


举报

相关推荐

百战c++(网络)

百战python04-循环结构

0 条评论