近年来,在生成图像建模中,生成对抗网络(GANs)的应用越来越多。基于风格的GAN可以生成不同层次的细节,从头部形状到眼睛颜色。它在高保真图像合成方面实现了SOTA,但其生成过程的计算复杂度非常高。难以应用于智能手机等移动设备。最近,一项专注于基于样式的生成模型性能优化的研究引起了人们的关注。这项研究分析了StyleGAN2中计算最困难的部分,并提出了对生成器网络的更改,使在边缘设备中部署基于样式的生成网络成为可能。该研究提出了一种名为MobileStyleGAN的新架构。与StyleGAN2相比,该架构的参数减少了约71%,计算复杂度降低了约90%,生成质量几乎没有损失。StyleGAN2(上)和MobileStyleGAN(下)生成的比较。该论文的作者在GitHub上放置了MobileStyleGAN的PyTorch实现。论文地址:https://arxiv.org/pdf/2104.04767.pdf项目地址:https://github.com/bes-dev/MobileStyleGAN.pytorch这个实现需要的训练代码很简单:StyleGAN2(左)和MobileStyleGAN(右)生成效果展示。让我们仔细看看MobileStyleGAN架构的方法细节。MobileStyleGAN架构MobileStyleGAN架构建立在基于样式的生成模型的基础上,包括映射网络和合成网络。前者使用StyleGAN2中的映射网络。本研究的重点是设计一个计算效率高的合成网络。MobileStyleGAN和StyleGAN2的区别StyleGAN2使用基于像素的图像表示,旨在直接预测输出图像的像素值。而MobileStyleGAN使用基于频率的图像表示,旨在预测输出图像的离散小波变换(DWT)。当应用于2D图像时,DWT将通道转换为四个具有较低空间分辨率和不同频带的大小相等的通道。然后,逆离散小波变换(IDWT)从小波域重建基于像素的表示,如下图所示。StyleGAN2利用跳跃生成器通过对同一图像的多个分辨率的RGB值进行显式求和来形成输出图像。研究发现,在小波域预测图像时,skipconnection-basedpredictionheads对生成图像的质量影响很小。因此,为了降低计算复杂度,本研究用网络中最后一个块的单个预测头替换了长跳生成器。但是从中间块预测目标图像对于稳定的图像合成具有重要意义。因此,本研究在每个中间块中加入一个辅助预测头,根据目标图像的空间分辨率对其进行预测。StyleGAN2和MobileStyleGAN的预测头差异。如下图所示,调制卷积由调制、卷积和归一化组成(左)。Depthwiseseparablemodulationconvolutions也包括这些部分(中)。StyleGAN2描述了权重的调制/解调,本研究将它们分别应用于输入/输出激活,这使得描述深度可分离调制卷积更容易。StyleGAN2构建块使用ConvTranspose(左下方)来放大输入特征图。相反,本研究使用IDWT作为MobileStyleGAN构建块(右下方)中的高档功能。由于IDWT不包含可训练参数,因此本研究在IDWT层之后添加了一个额外的深度可分离调制卷积。StyleGAN2和MobileStyleGAN完整的buildingblock结构如下图所示:基于蒸馏的训练过程与之前的一些研究类似,本研究的训练框架也是基于知识蒸馏技术。该研究使用StyleGAN2作为教师网络来训练MobileStyleGAN以模仿StyleGAN2的功能。训练框架如下图所示。
