「百战GAN」StyleGAN人脸属性(表情年龄性别)编辑代码实战

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

作者&编辑 | 言有三

本文篇幅：7000字

背景要求：会使用Python和Pytorch

附带资料：参考论文和项目，视频讲解

1 项目背景

作为一门新兴的技术，GAN在很多人脸图像任务中有着广泛的应用，下面列举了一些典型的方向。

比如使用SRGAN进行人脸超分辨，使用BeautyGAN进行人脸美颜，使用Deepfake进行换脸。如今，基于人脸属性编辑的技术在抖音快手等应用中都是非常受欢迎的技术，比如可以体验让人变老变小，如下图年龄的更改。

这背后的核心技术之一就是在生成模型的Latent空间进行属性向量的编辑，然后通过生成模型来获取最终结果图，其中最典型的代表性框架就是StyleGAN，本次我们来实战使用StyleGAN进行人脸属性编辑。

在阅读接下来的内容之前，请大家务必回顾之前StyleGAN人脸生成的内容，没有这一部分内容基础，无法学习接下来的内容。

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

【项目实战课】基于Pytorch的StyleGAN v1人脸图像生成实战

2 原理简介

现在我们要使用StyleGAN进行真实人脸的编辑，需要解决2个关键问题。

(1) 如何获得真实人脸的潜在编码向量，它对应StyleGAN中的映射网络（mapping network）的输入Z或者输出W。

(2) 如何通过修改Z或者W向量，控制生成人脸图像的高层语义属性。

接下来我们重点介绍潜在编码向量的求解，并在下一节中实验基于潜在编码向量的属性编辑。

当前对真实人脸编码向量的求取基本上基于两种思路，一种是学习一个编码器来实现映射，一种是直接对向量进行优化求解。

2.1 基于编码器的求解

基于编码器的求解框架如下图所示，由两部分模块组成。

Encoder表示需要训练的编码器，Decoder表示已经训练完的生成模型，如StyleGAN的生成器部分。真实图像输入编码器得到Z或者W，再输入生成器得到生成的人脸，完成人脸图像的重建。

通过直接学习一个编码器，可以不需要对每一张图都进行优化，实现一次训练，对任意图像都能提取潜在编码向量，但是也容易在训练数据集上发生过拟合。

2.2 基于优化求解的方法

另外一种方法就是基于优化求解，直接对每一张图片优化出对应的W，在StyleGAN v2，Image2StyleGAN等框架中采用了这种方案，并且512维的W被拓展成W+, W+为18×512维的矩阵，这样就可以对每一个自适应实例归一化（AdaIN）风格模块都使用不同W，实现更自由的属性编辑。

基于优化求解的方法包括以下几步：

(1) 给定图片I，以及预训练好的生成器G。

(2) 初始化潜在编码向量，如W，其初始值可以使用计算得到的统计平均值。

(3) 根据优化目标进行反复迭代，直到达到预设的终止条件。

优化目标的常见形式为：

其中Lpercept为特征空间中的感知损失距离，是很通用的问题，具体形式不再赘述。lamda用于平衡感知损失和MSE损失之间的权重比。可以使用梯度下降算法进行求解。

基于优化求解方法的优点是精度较高，但是优化速度慢，而且对每一个图片都必须进行优化迭代。

在求解得到了潜在编码向量后，我们就可以通过编辑向量来编辑人脸的高层语义属性，对于StyleGAN架构来说，潜在编码向量可以是Z也可以是W，一般基于W进行编码会有更好的效果。

3 人脸属性编辑实战

接下来我们实践基于StyleGAN模型的人脸属性编辑。

3.1 人脸重建

要使用StyleGAN来进行人脸编辑，首先我们需要将人脸投射到潜在编码向量空间，下面我们采用的方法是基于优化的方法，即对每一张人脸图片，单独优化求解出潜在编码向量，基本思想前面已经介绍过。

接下来我们来看人脸重建的求解，核心代码如下：

if __name__ == "__main__":
    ## 预训练模型权重
    parser.add_argument(
        "--ckpt", type=str, required=True, help="path to the model checkpoint"
    )

    ## 输出图像尺寸
    parser.add_argument(
        "--size", type=int, default=256, help="output image sizes of the generator"
    )

    ## 学习率参数
    parser.add_argument(
        "--lr_rampup",
        type=float,
        default=0.05,
        help="duration of the learning rate warmup",
    )
    parser.add_argument(
        "--lr_rampdown",
        type=float,
        default=0.25,
        help="duration of the learning rate decay",
    )
    parser.add_argument("--lr", type=float, default=0.1, help="learning rate")


    ## 噪声相关参数，噪声水平，噪声衰减范围，噪声正则化
    parser.add_argument(
        "--noise", type=float, default=0.05, help="strength of the noise level"
    )
    parser.add_argument(
        "--noise_ramp",
        type=float,
        default=1.0,
        help="duration of the noise level decay",
    )
    parser.add_argument(
        "--noise_regularize",
        type=float,
        default=10000,
        help="weight of the noise regularization",
    )

    ## MSE损失
    parser.add_argument("--mse", type=float, default=0.5, help="weight of the mse loss")

    ## 迭代次数
    parser.add_argument("--step", type=int, default=1000, help="optimize iterations")

    ## 重建图像
    parser.add_argument(
        "--files", type=str, help="path to image files to be projected"
    )

    ## 重建结果
    parser.add_argument(
        "--results", type=str, help="path to results files to be stored"
    )

## 计算学习率
def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
    lr_ramp = min(1, (1 - t) / rampdown)
    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
lr_ramp = lr_ramp * min(1, t / rampup)

return initial_lr * lr_ramp

## 合并latent向量和噪声
def latent_noise(latent, strength):
noise = torch.randn_like(latent) * strength

    return latent + noise

## 生成图片
def make_image(tensor):
    return (
        tensor.detach()
        .clamp_(min=-1, max=1)
        .add(1)
        .div_(2)
        .mul(255)
        .type(torch.uint8)
        .permute(0, 2, 3, 1)
        .to("cpu")
        .numpy()
    )

## 生成与图像大小相等的噪声
def make_noise(device,size):
    noises = []
    step = int(math.log(size, 2)) - 2
    for i in range(step + 1):
            size = 4 * 2 ** i
            noises.append(torch.randn(1, 1, size, size, device=device))
return noises

## 噪声归一化
def noise_normalize_(noises):
    for noise in noises:
        mean = noise.mean()
        std = noise.std()
        noise.data.add_(-mean).div_(std)

args = parser.parse_args()

    device = "cpu"
    ## 计算latent向量的平均次数
    n_mean_latent = 10000

    ## 获得用于计算损失的最小图像尺寸
    resize = min(args.size, 256)

    ## 预处理函数
    transform = transforms.Compose(
        [
            transforms.Resize(resize),
            transforms.CenterCrop(resize),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        ]
    )

    ## 投影的人脸图片,将图片处理成一个batch
    imgs = []
    imgfiles = os.listdir(args.files)
    for imgfile in imgfiles:
        img = transform(Image.open(os.path.join(args.files,imgfile)).convert("RGB"))
        imgs.append(img)

    imgs = torch.stack(imgs, 0).to(device)

    ## 载入模型
    netG = StyledGenerator(512,8)
    netG.load_state_dict(torch.load(args.ckpt,map_location=device)["g_running"], strict=False)
    netG.eval()
    netG = netG.to(device)
    step = int(math.log(args.size, 2)) - 2
    with torch.no_grad():
        noise_sample = torch.randn(n_mean_latent, 512, device=device)
        latent_out = netG.style(noise_sample) ##输入噪声向量Z，输出latent向量W=latent_out
        latent_mean = latent_out.mean(0)
        latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5

    ## 感知损失计算
    percept = lpips.PerceptualLoss(
        model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
    )

    ## 构建噪声输入
    noises_single = make_noise(device,args.size)

    noises = []
    for noise in noises_single:
        noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())

    ## 初始化Z向量
    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
    latent_in.requires_grad = True

    for noise in noises:
        noise.requires_grad = True

    optimizer = optim.Adam([latent_in] + noises, lr=args.lr)

    pbar = tqdm(range(args.step))

    ## 优化学习Z向量
    for i in pbar:
        t = i / args.step ## t的范围是(0,1)
        lr = get_lr(t, args.lr)
        optimizer.param_groups[0]["lr"] = lr

## 噪声衰减
        noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2
        latent_n = latent_noise(latent_in, noise_strength.item())
        latent_n.to(device)
        img_gen = netG([latent_n], noise=noises, step=step) ## 生成的图片
        batch, channel, height, width = img_gen.shape

## 在不超过256的分辨率上计算损失
        if height > 256:
            factor = height // 256

            img_gen = img_gen.reshape(
                batch, channel, height // factor, factor, width // factor, factor
            )
            img_gen = img_gen.mean([3, 5])

        p_loss = percept(img_gen, imgs).sum()  ## 感知损失
        n_loss = noise_regularize(noises)      ## 噪声损失
        mse_loss = F.mse_loss(img_gen, imgs)   ## MSE损失

        loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        noise_normalize_(noises)

        pbar.set_description(
            (
                f"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f};"
                f" mse: {mse_loss.item():.4f}; loss: {loss.item():.4f}; lr: {lr:.4f}"
            )
        )

    ## 重新生成高分辨率图片
    img_gen = netG([latent_in], noise=noises,step=step)
    img_ar = make_image(img_gen)

    result_file = {}
    for i, input_name in enumerate(imgfiles):
        noise_single = []
        for noise in noises:
            noise_single.append(noise[i : i + 1])

        print("i="+str(i)+"; len of imgs:"+str(len(img_gen)))
        result_file[input_name] = {
            "img": img_gen[i],
            "latent": latent_in[i],
            "noise": noise_single,
        }

        img_name = os.path.join(args.results,input_name)
        pil_img = Image.fromarray(img_ar[i])
        pil_img.save(img_name) ##存储图片
        np.save(os.path.join(args.results,input_name.split('.')[0]+'.npy'),latent_in[i].detach().numpy())##存储latent向量

在上面的代码中，latent_in就是要优化学习的潜在变量，当使用netG([latent_n], noise=noises, step=step)的调用方式时，latent_n是映射网络（mapping network）的向量Z，它由latent_in和随着迭代不断衰减的噪声向量组成，此时没有输入平均风格向量。

根据之前我们对生成器模型结构的解读，此时只有latent_n会影响生成的风格，因此为基于向量Z的重建方法。

当然我们也可以通过将latent_n设为平均风格向量，将混合权重style_weight设置为0，这样也只有latent_n会影响生成的风格，并且它会作为生成网络（synthetis network）的输入，即向量W，此时就是基于向量W的重建方法。

关于感知损失的计算，以及预训练模型的解读，请大家自行完成，或者阅读文后的讲解视频。

下面我们会比较这两种重建方法的差异。

学习率采用warmup策略,即先增大再减小，最大不超过0.1。感知损失、MSE损失、噪声正则化损失的权重分别为1.0，1.0,10000,噪声相关的幅度、衰减范围因子分别为0.05，1。

下图展示了随机选取的图像的人脸重建结果图。

上图中第一行是原图，第二行是基于向量Z的重建图，第三行是基于向量W的重建图。

下图展示了一幅图像基于Z和基于W向量的3部分训练损失曲线图，其中实线对应W，虚线对应Z，需要注意的是噪声损失没有乘以对应的权重，如果乘以对应的权重，最终收敛的时候三部分损失的幅度相当。

从结果上来看，两种方法中人脸图像的总体姿态、肤色、发型、脸型、背景重建效果都比较好，基于向量Z的重建方法中人脸的清晰度更高，但是身份没有得到保持。这主要是因为所学习的特征向量为Z，它还需要经过非线性的映射网络（mapping network）得到W，相比于直接学习W的难度更高，许多研究都表明使用W向量可以得到更好的重建效果。

从损失曲线可以看出，基于W向量的重建训练损失可以获得更低的值，但是感知损失的收敛更慢，MSE的实际损失更低，可见在获得了更精确的身份重建的基础上，牺牲了一定的感知质量，使得当编辑W向量时会比较敏感。

接下来我们使用基于向量Z的重建结果来进行人脸的属性编辑，因为这样可以比较基于Z和W来编辑属性的差异性，基于向量W的重建结果中无法获得对应的Z向量，因为映射网络是无法从输出获得输入的。

3.2 人脸属性混合与插值

接下来我们进行人脸属性的混合与插值，这是在多张图片之间进行属性的混合操作。

(1) 人脸属性样式混合

我们首先体验人脸属性样式混合，样式混合操作可以通过下式的向量运算实现。

其中[0:m]表示取向量的前m维，[m:n]表示取向量的第m到n维，两者通过拼接得到新的向量。

值得注意的是，这里的W并不是映射网络的输出向量，而是对于不同风格化层的AdaIN缩放和偏移系数，即风格系数。我们之前介绍过，不同分辨率的风格模块对应着不同层级的人脸特征，在这里我们使用两张人脸图像对应的风格系数进行混合，体验不同层级特征的样式混合。

图中第1列表示源图，第5列表示目标图，第2列、第3列、第4列表示在对应分辨率的风格层使用源图的风格，其他分辨率的风格层使用目标图的风格。

第2列表示分辨率为4×4，8×8的风格化模块的风格向量来自于源图，其他分辨率模块的风格向量来自于目标图。可以看出，结果图有源图的粗粒度特征，如人脸姿态、发型特征，以及目标图的细粒度特征，如发丝、眼睛颜色。

第3列表示分辨率为16×16，32×32的风格化模块的风格向量来自于源图，其他分辨率模块的风格向量来自于目标图。可以看出，保留了源图的中等粒度特征，如眼睛形态、嘴唇颜色等特征。

第4列表示分辨率为64×64到1024×1024之间的风格化模块的风格向量来自于源图，其他分辨率模块的风格向量来自于目标图。可以看出，结果图有源图的细粒度特征，如头发和皮肤的颜色与纹理等特征，以及目标图的粗粒度特征，如姿态、发型。

(2) 人脸样式插值

接下来我们体验人脸样式插值，样式插值操作可以通过下面的W向量运算实现。

下图分别展示了基于向量Z和向量W的人脸样式插值结果。

第1列表示A域图像，对应n维的列向量 。第6列表示B域图像，对应n维的列向量，第2，3，4，5列表示在不同权重下对两张图片的向量进行加权后生成的图像，可以看出都可以实现了样式过度，但是基于W向量的结果明显要比基于Z向量的结果更加平滑。

3.2 人脸属性编辑

上面介绍的人脸样式混合可以直接通过两幅图片的潜在向量运算来实现人脸属性的混合，而如果想要对单张人脸的属性进行精确编辑，就需要首先找到潜在向量的编辑方向，称之为方向向量。下面我们介绍方向向量的求解与基于方向向量的属性编辑。

接下来的属性编辑建立在以下假设的基础上：在方向向量上，线性的改变潜在编码向量，则生成的图像及语义内容也是连续变化的，因此可以使用线性模型来进行属性更改：

W表示结果编码，W0表示人脸特征码，alpha表示偏移系数，n表示方向向量。

接下来需要解决的问题就是方向向量n的求解，具体的步骤为：

(1) 随机采样潜在编码向量，生成人脸图片，保存人脸图片和对应的潜在编码向量。

(2) 对生成的人脸图片，训练想要编辑的人脸属性CNN分类模型。任意一个二值语义都存在一个超平面可以作为语义类别的分类边界，在超平面的一侧改变潜在编码向量不会改变对应的语义类别，这个超平面可以用单位法线矢量表示。

(3) 根据CNN分类模型获得的标签，对潜在编码向量训练出线性回归模型，获得方向向量，如下图所示。

上图展示了0和1这两类样本，它就是潜在编码向量W。求解完线性回归模型的权重 ，它实际上就是方向向量。得到方向向量后，就可以进行人脸相关属性的编辑。

以人脸表情编辑为例：首先我们使用StyleGAN随机生成了50000张人脸图片，接下来我们需要将图片分为有表情和无表情两类，需要使用一个预训练好的2分类表情识别模型。为了保证模型有较高的准确率，训练时采用的方法是首先提取出人脸嘴唇区域，然后对嘴唇区域进行训练。

该表情识别模型大家可以参考阅读：

【项目实战课】AI零基础，人人免费可学！基于Pytorch的SimpleNet人脸表情识别实战

下图就是基于微笑表情模型分类出的有微笑和无微笑两类图片的一些样本展示。

训练好模型后，我们就得到了图片的标签，然后对其对应的潜在向量学习一个线性分类器，得到方向向量。完整代码请大家参考文后资料。

因为潜在向量维度是1×512，所以得到的方向向量也是1×512维，然后我们就可以基于方向向量进行属性编辑，基于Z向量编辑的核心代码如下：

from model import StyledGenerator
if __name__ == "__main__":
    device = "cpu"
    parser.add_argument(
        "--ckpt", type=str, required=True, help="path to the model checkpoint"
    )
    parser.add_argument(
        "--size", type=int, default=1024, help="output image sizes of the generator"
    )
    parser.add_argument(
        "--files", type=str, help="path to image files to be projected"
    )
    parser.add_argument(
        "--direction", type=str, help="direction file to be read"
    )
    parser.add_argument(
        "--directionscale", type=float, help="direction scale"
    )
    args = parser.parse_args()

    ## 载入模型
    netG = StyledGenerator(512)
    netG.load_state_dict(torch.load(args.ckpt,map_location=device)["g_running"], strict=False)
    netG.eval()
    netG = netG.to(device)
    step = int(math.log(args.size, 2)) - 2

    ## 载入方向向量
    direction = np.load(args.direction)
    directiontype = args.direction.split('/')[-1].split('.')[0]
    editscale = args.directionscale
    npys = glob.glob(args.files+"*.npy")

    for npyfile in npys:
        latent = torch.from_numpy(np.load(npyfile))
        if len(latent.shape) == 1:
            latent = latent.unsqueeze(0)
        latent = latent + torch.from_numpy((editscale*direction[0]).astype(np.float32))
        latent.to(device)
        img_gen = netG([latent], step=step) ##生成的图片
        img_name = os.path.join(npyfile.replace('.npy','_'+directiontype+'_'+str(editscale)+'.jpg'))
        utils.save_image(img_gen, img_name, normalize=True)
        np.save(img_name.replace('.jpg','.npy'),latent)

下图分别展示了基于Z向量和W向量的编辑结果。

第1行表示原图，第2行表示减小微笑表情幅度，第3行表示增大微笑表情幅度。

可以看出基于Z向量和W向量，都能够实现微笑表情的编辑。不过基于Z向量的模型，明显地更改了其他的属性，比如发型、人脸身份等，这说明基于向量Z的编辑不能很好地实现表情属性与其他属性的解耦合，存在非常大的改进空间，而基于W向量能够比较好的保护其他属性信息。

由于不同的人脸有不同的属性，我们也可以基于此进行属性的添加与移除，W向量运算如下式，

其中W3和W2来自于同一个人，相减得到某个属性的向量，然后将其添加到W1，表示给对应的人脸添加对应的属性。

下图分别展示了基于Z向量和W向量的人脸微笑属性添加与移除结果。

第1列表示源图，对应n维的列向量 。第2列和第3列分别表示目标图，对应n维的列向量和。第4列表示生成的图像，可以看出基于Z向量虽然可以一定程度上实现微笑属性的添加与移除，但是会修改人脸的身份。而基于W向量可以很好地实现微笑属性的添加与移除。

本文参考的文献如下：

[1] Abdal R , Qin Y , Wonka P . Image2StyleGAN: How to Embed Images Into the StyleGAN Latent Space?[J]. IEEE, 2019.

[2] Shen Y , Gu J , Tang X , et al. Interpreting the Latent Space of GANs for Semantic Face Editing[C]// 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2020.

本文视频讲解和代码，请大家移步：【项目实战课】基于Pytorch的StyleGAN人脸属性（表情、年龄、性别）编辑实战

总结

本次我们使用StyleGAN模型实践了人脸属性的编辑，取得了预期的实验结果，不过该方法还存在许多可以改进的地方，大家可以使用更新的StyleGAN模型，以及更新的属性编辑模型来进行改进，欢迎大家以后持续关注《百战GAN专栏》。

如何系统性地学习生成对抗网络GAN

欢迎大家关注有三AI-CV秋季划GAN小组，可以系统性学习GAN相关的内容，包括GAN的基础理论，《深度学习之图像生成GAN：理论与实践篇》，《深度学习之图像翻译GAN：理论与实践篇》以及各类GAN任务的实战。

介绍如下：【CV秋季划】生成对抗网络GAN有哪些研究和应用，如何循序渐进地学习好（2022年言有三一对一辅导）？