Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

小舟、陈萍作者

压缩版styleGAN,合成高保真图像,参数更少、计算复杂度更低

一个名为 MobileStyleGAN 的新架构大大减少了基于样式 GAN 的参数量,降低了计算复杂度。

近年来在生成图像建模中,生成对抗网络(GAN)的应用越来越多。基于样式(style-based)的 GAN 可以生成不同层次的细节,大到头部形状、小到眼睛颜色,它在高保真图像合成方面实现了 SOTA,但其生成过程的计算复杂度却非常高,难以应用于智能手机等移动设备。


近日,一项专注于基于样式的生成模型的性能优化的研究引发了大家的关注。该研究分析了 StyleGAN2 中最困难的计算部分,并对生成器网络提出了更改,使得在边缘设备中部署基于样式的生成网络成为可能。该研究提出了一种名为 MobileStyleGAN 的新架构。相比于 StyleGAN2,该架构的参数量减少了约 71 %,计算复杂度降低约 90 %,并且生成质量几乎没有下降。

StyleGAN2(上)与 MobileStyleGAN(下)的生成效果对比。

论文作者已将 MobileStyleGAN 的 PyTorch 实现放到了 GitHub 上。


  • 论文地址:https://arxiv.org/pdf/2104.04767.pdf

  • 项目地址:https://github.com/bes-dev/MobileStyleGAN.pytorch


该实现所需的训练代码非常简单:


pip install -r requirements.txtpython train.py --cfg configs/mobile_stylegan_ffhq.json --gpus <n_gpus>

StyleGAN2(左)与 MobileStyleGAN(右)的生成效果展示。

下面我们来具体看一下 MobileStyleGAN 架构的方法细节。

MobileStyleGAN 架构


MobileStyleGAN 架构是在基于样式生成模型的基础上构建的,它包括映射网络和合成网络,前者采用的是 StyleGAN2 中的映射网络,该研究的重点是设计了一个计算高效的合成网络。


MobileStyleGAN 与 StyleGAN2 的区别

StyleGAN2 使用基于像素的图像表征,并旨在直接预测输出图像的像素值。而 MobileStyleGAN 使用基于频率的图像表征,旨在预测输出图像的离散小波变换 (DWT)。当应用到 2D 图像,DWT 将信道转换成四个大小相同的信道,这几个信道具有较低的空间分辨率和不同的频带。然后,逆向离散小波变换(IDWT) 从小波域重建基于像素的表征,如下图所示。


StyleGAN2 利用跳远生成器(skip-generator),通过对同一图像的多个分辨率的 RGB 值进行显式求和来形成输出图像。该研究发现,当在小波域中对图像进行预测时,基于跳远连接(skip connection)的预测头对生成图像的质量影响不大。因此,为了降低计算复杂度,该研究采用网络中最后一个块的单个预测头替换跳远生成器。但从中间块中预测目标图像对于稳定的图像合成具有重要意义。因此,该研究为每个中间块添加一个辅助预测头,根据目标图像的空间分辨率对其进行预测。

StyleGAN2 和 MobileStyleGAN 的预测头区别。

如下图所示,调制卷积包括调制、卷积和归一化(左)。深度可分离调制卷积也包括这些部分(中)。StyleGAN2 描述了用于权重的调制 / 解调,该研究分别将它们应用于输入 / 输出激活,这使得描述深度可分离调制卷积更加容易。



StyleGAN2 构造块使用 ConvTranspose(下图左)来 upscale 输入特征映射。而该研究在 MobileStyleGAN 构造块(下图右)中使用 IDWT 当作 upscale 函数。由于 IDWT 不包含可训练参数,该研究在 IDWT 层之后增加了额外的深度可分离调制卷积。


StyleGAN2 和 MobileStyleGAN 的完整构造块结构如下图所示:



基于蒸馏的训练过程

类似于此前的一些研究,该研究的训练框架也基于知识蒸馏技术。该研究将 StyleGAN2 作为教师网络,训练 MobileStyleGAN 来模仿 StyleGAN2 的功能,训练框架如下图所示。

理论StyleGAN
2
相关数据
权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

映射技术

映射指的是具有某种特殊结构的函数,或泛指类函数思想的范畴论中的态射。 逻辑和图论中也有一些不太常规的用法。其数学定义为:两个非空集合A与B间存在着对应关系f,而且对于A中的每一个元素x,B中总有有唯一的一个元素y与它对应,就这种对应为从A到B的映射,记作f:A→B。其中,y称为元素x在映射f下的象,记作:y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域,记作f(A)。同样的,在机器学习中,映射就是输入与输出之间的对应关系。

生成模型技术

在概率统计理论中, 生成模型是指能够随机生成观测数据的模型,尤其是在给定某些隐含参数的条件下。 它给观测值和标注数据序列指定一个联合概率分布。 在机器学习中,生成模型可以用来直接对数据建模(例如根据某个变量的概率密度函数进行数据采样),也可以用来建立变量间的条件概率分布。

生成对抗网络技术

生成对抗网络是一种无监督学习方法,是一种通过用对抗网络来训练生成模型的架构。它由两个网络组成:用来拟合数据分布的生成网络G,和用来判断输入是否“真实”的判别网络D。在训练过程中,生成网络-G通过接受一个随机的噪声来尽量模仿训练集中的真实图片去“欺骗”D,而D则尽可能的分辨真实数据和生成网络的输出,从而形成两个网络的博弈过程。理想的情况下,博弈的结果会得到一个可以“以假乱真”的生成模型。

生成对抗技术

生成对抗是训练生成对抗网络时,两个神经网络相互博弈的过程。两个网络相互对抗、不断调整参数,最终目的是使判别网络无法判断生成网络的输出结果是否真实。

推荐文章
暂无评论
暂无评论~