仅17 KB、一万个权重的微型风格迁移网络！

点击上方蓝色字体，关注我们

来源：机器之心译

今天 reddit 上一篇帖子引起了热议，博主 jamesonatfritz 称他将原本具备 1.7M 参数的风格迁移网络减少到只有 11,868 个参数，该网络仍然能够输出风格化的图像。且量化后的最终网络体积仅有 17 kB，非常适合移动 app。

jamesonatfritz 想解决 神经网络 的过 参数 化问题，想要创建体积小但性能优的 神经网络 。他所试验的第一个任务便是艺术风格迁移。

• GitHub 链接：

• https://github.com/fritzlabs/fritz-style-transfer

现在有很多用来训练艺术风格迁移模型的现成工具，还有上千种开源实现。其中的多数工具利用 Johnson 等人在《Perceptual Losses for Real-Time Style Transfer and Super-Resolution》中提出的网络架构的变体来实现快速、前馈的风格化。因此，多数迁移模型的大小是 7MB。对于你的应用来说，这个负担并非不可承受，但也并非无足轻重。

研究表明， 神经网络 的体积通常远远大于所需，数百万的 权重 中有很多并不重要。因此作者创造了一个体积大大缩小的可靠风格迁移模型：一个只有 11686 个训练 权重 的 17KB 神经网络 。

左：原图；中：来自上述 17KB 模型的风格化图像；右：来自 7MB 模型的风格化图像。

快速概览：

原始模型：

• 大小：7MB

• 权重 数：1.7M

• 在 iPhone X 上的速度：18 FPS

小模型：

• 大小：17KB

• 权重 数：11,868

• 在 iPhone X 上的速度：29 FPS

如何缩小风格迁移模型

作者主要使用了两种技术，而且都可以泛化到其他模型：

1. 大刀阔斧地修剪层和 权重 ；

2. 通过量化将 32 位浮点 权重 转换为 8 位整型

修剪策略

卷积 神经网络 通常包含数百万甚至上亿个需要在训练阶段进行调整的 权重 。通常来讲， 权重 越多 准确率 越高。但这种增加 权重 提高 准确率 的做法非常低效。谷歌 MobileNet V2 的 stock 配置具有 347 万个 权重 ，内存占用达 16MB。InceptionV3 架构大小约为前者的 6 倍，具备 2400 万个 权重 ，内存占用达 92MB。尽管多了 2000 多万个 权重 ，但 InceptionV3 在 ImageNet 上的 top-1 分类 准确率 只比 MobileNet V2 高出 7 个百分点（80% vs 73%）。

因此，我们可以假设 神经网络 中的多数 权重 没有那么重要并将其移除。但重点是怎么做呢？我们可以选择在三个层面进行修剪：单个 权重 、层、块。

权重 层面：假设某个 神经网络 上的多数（>95%） 权重 都没有什么用。如果能找出那些对 准确率 有影响的 权重 ，就可以将其留下并将其他移除。

层层面：每个层中都包含一些 权重 。例如，2D 卷积层具有一个 权重 张量 ，即卷积核，用户可以定义其宽度、高度和深度。缩小卷积核可以减小整个网络的大小。

块层面：多个层通常可以结合成可重复利用的子图，即块。以 ResNet 为例，它的名字来源于重复 10-50 次的「残差块」。在块层面进行修剪可以移除多个层，从而一次性移除多个 参数 。

在实践中，稀疏 张量 运算没有很好的实现，因此 权重 层面的修剪没有多大价值。那么就只剩下层和块层面的修剪了。

实践中的修剪

作者使用的层修剪技术是引入 width multiplier 作为超 参数 。width multiplier 最初由谷歌在其论文《 MobileNet s: Efficient Convolutional Neural Network s for Mobile Vision》中提出，非常简单、高效。

width multiplier 利用一个恒定系数调整每个卷积层中的卷积核数量。对于给定的层及 width multiplier alpha，卷积核数量 F 变为 alpha * F。

有了这个超 参数 ，我们就可以生成一系列架构相同但 权重 数不同的网络。训练每种配置，就可以在模型速度、大小及 准确率 之间做出权衡。

下面是作者模仿 Johnson 等人在《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Network s》提出的网络架构构建快速风格迁移模型的方法，不同之处在于添加了 width multiplier 作为超 参数 。

@classmethod
def build(
        cls,
        image_size,
        alpha=1.0,
        input_tensor=None,
        checkpoint_file=None):
    """Build a Transfer Network Model using keras' functional API.
    Args:
        image_size - the size of the input and output image (H, W)
        alpha - a width parameter to scale the number of channels by
    Returns:
        model: a keras model object
    """
    x = keras.layers.Input(
        shape=(image_size[0], image_size[1], 3), tensor=input_tensor)
    out = cls._convolution(x, int(alpha * 32), 9, strides=1)
    out = cls._convolution(out, int(alpha * 64), 3, strides=2)
    out = cls._convolution(out, int(alpha * 128), 3, strides=2)
    out = cls._residual_block(out, int(alpha * 128))
    out = cls._residual_block(out, int(alpha * 128))
    out = cls._residual_block(out, int(alpha * 128))
    out = cls._residual_block(out, int(alpha * 128))
    out = cls._residual_block(out, int(alpha * 128))
    out = cls._upsample(out, int(alpha * 64), 3)
    out = cls._upsample(out, int(alpha * 32), 3)
    out = cls._convolution(out, 3, 9, relu=False, padding='same')
    # Restrict outputs of pixel values to -1 and 1.
    out = keras.layers.Activation('tanh')(out)
    # Deprocess the image into valid image data. Note we'll need to define
    # a custom layer for this in Core ML as well.
    out = layers.DeprocessStylizedImage()(out)
    model = keras.models.Model(inputs=x, outputs=out)

注意，模型构建器类的其余部分没有显示。

当 alpha=1.0 时，得到的网络包含 170 万个 权重 。当 alpha=0.5 时，得到的网络仅有 424,102 个 权重 。

你可以构建一些宽度 参数 很小的网络，但是也有相当多的重复块。作者决定修剪掉一些，但实际操作后却发现不能移除太多。即使 参数 量保持不变，较深的网络能够产生更好的结果。作者最终删除了五个残差块中的两个，并将每层的默认滤波器数量减少至 32 个。得到的微型网络如下所示：

@classmethod
def build(
        cls,
        image_size,
        alpha=1.0,
        input_tensor=None,
        checkpoint_file=None):
    """Build a Small Transfer Network Model using keras' functional API.
    This architecture removes some blocks of layers and reduces the size
    of convolutions to save on computation.
    Args:
        image_size - the size of the input and output image (H, W)
        alpha - a width parameter




    
 to scale the number of channels by
    Returns:
        model: a keras model object
    """
    x = keras.layers.Input(
        shape=(image_size[0], image_size[1], 3), tensor=input_tensor)
    out = cls._convolution(x, int(alpha * 32), 9, strides=1)
    out = cls._convolution(out, int(alpha * 32), 3, strides=2)
    out = cls._convolution(out, int(alpha * 32), 3, strides=2)
    out = cls._residual_block(out, int(alpha * 32))
    out = cls._residual_block(out, int(alpha * 32))
    out = cls._residual_block(out, int(alpha * 32))
    out = cls._upsample(out, int(alpha * 32), 3)
    out = cls._upsample(out, int(alpha * 32), 3)
    out = cls._convolution(out, 3, 9, relu=False, padding='same')
    # Restrict outputs of pixel values to -1 and 1.
    out = keras.layers.Activation('tanh')(out)
    # Deprocess the image into valid image data. Note we'll need to define
    # a custom layer for this in Core ML as well.
    out = layers.DeprocessStylizedImage()(out)
    model = keras.models.Model(inputs=x, outputs=out)

带有宽度 参数 的较小风格迁移网络。

通过反复尝试，作者发现仍然可以用上述架构实现良好的风格迁移，一直到宽度 参数 为 0.3，在每一层上留下 9 个滤波器。最终结果是一个只有 11,868 个 权重 的 神经网络 。任何 权重 低于 10000 的网络都不能持续训练，并且会产生糟糕的风格化图像。

值得一提的是， 剪枝 技术是在网络训练之前应用的。在训练期间或训练后反复修剪，你可以在很多任务上实现更高的性能。

量化

最后一段压缩是在网络训练完成后进行的。 神经网络 权重 通常存储为 64 位或 32 位浮点数。量化过程将每一个浮点 权重 映射 到具有较低位宽的整数。从 32 位浮点 权重 变为 8 位整型，使得存储大小减少了 4 倍。作者利用 Alexis Creuzot 在博客中提出的方法（https://heartbeat.fritz.ai/reducing-coreml2-model-size-by-4x-with-quantization-in-ios12-b1c854651c4），在不怎么影响风格的情况下使浮点数降低到了 8 位量化。

现在所有主要的移动框架都支持量化，如 Tensor Flow Mobile、 Tensor Flow Lite、Core ML 和 Caffe2Go。

最终结果

该微型网络架构有 11,868 个 参数 ，相比之下，Johnson 最初的模型具有 170 万个 参数 ，大小为 1.7MB。当转化为 Core ML 并量化时，最终大小仅为 17KB——为原始大小的 1/400。以下是在梵高的《Starry Night》上的训练结果。

此微型风格迁移结果的实时视频可在 Heartbeat App 上查看：

http://bit.ly/heartbeat-ios

作者惊讶地发现，尽管尺寸相差 400 倍，但在 iPhone X 上，这款微型模型的运行速度仅快了 50%。原因可能是计算与这一通用架构相关，也可能是将图像迁移到 GPU 进行处理时造成的。

如果你对结果表示怀疑，可以自己下载并运行此微型模型。甚至训练自己的模型！

• 下载地址：

• https://github.com/fritzlabs/fritz-style-transfer/blob/master/example/starry_night_640x480_small_a03_q8.mlmodel

总而言之，作者用两种简单的技术将风格迁移 神经网络 的规模减小了 99.75%。使用简单的 width multiplier 超 参数 修剪层，训练后的 权重 从 32 位浮点数量化为 8 位整数。未来，作者期待看到将这些方法泛化到其它 神经网络 的效果。风格迁移相对简单，因为「 准确率 」肉眼可见。对于图像识别这样更加可以量化的任务而言，如此极端的修剪可能带来更明显的性能下降。

Reddit 讨论

这篇帖子下有一些 reddit 网友对该项目提出了质疑：

gwern：

看你的博客， 剪枝 部分似乎没有移除任何层，只是更改了层的宽度／滤波器，然后对所有 参数 进行量化。如果所有层都在（因为你没有做任何类似于训练较宽的浅层网络的工作来模仿原始深度教师网络），那么它们仍将从后续计算的每一层中引入大量延迟，即使每一层都很小。（由于你可以在手机 GPU 上安装更多模型，每个模型使用较少的 FLOPS，因此整体吞吐量会变得更好。但是每个模型的迭代仍然需要一段时间，在特定大小之后，每一层基本上是即时的。）

jamesonatfritz 回复：

你说得对，滤波器 剪枝 部分确实没有移除层，但是我去掉了两个残差块，从而消除了一些层。整体 FLOPs 的降低情况不如全部 权重 数量的减少情况，这一点你说得对。不幸的是，Apple 没法让你较好地控制模型运行的位置。你无法强制该模型使用 GPU。一些启发式方法导致较小的模型仅在 CPU 上运行，这是可能的。

gwern：

「去掉了两个残差块，从而消除了一些层。」

你借此获得了一些加速，但是使用更扁平的模型或许会实现更多加速。

jamesonatfritz 回复：

确实如此。我试过的最扁平模型只有一个卷积层、一个残差模块和一个上采样模块，但我发现这些变体无法 收敛 。

gwern：

这似乎有些过了：只有一个层有些过于难了。我想的是三四个层这样，在预训练风格迁移模型的确切像素输出上进行训练。或许值得一试。

Ikuyas：

这个方法和直接使用小模型有什么区别吗？我确定使用 11,868 个 参数 进行训练结果会更好。另外，1.7M 参数 太大了，每个 参数