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

选自fritz.ai，机器之心编译，机器之心编辑部。

今天 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 位整型

修剪策略

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

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

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

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

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

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

实践中的修剪

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

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

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

下面是作者模仿 Johnson 等人在《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》提出的网络架构构建快速风格迁移模型的方法，不同之处在于添加了 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 位量化。

现在所有主要的移动框架都支持量化，如 TensorFlow Mobile、TensorFlow 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 参数太大了，每个参数的贡献估计会很小。但是，真实情况是只有几百个参数是真正重要的，其他参数只是到处吸收一点微小的噪声。从标准回归的角度来看，这似乎是完美的预期结果。

作者回复：

我应该在文章里写清楚的，事实上你所说的正是我所做的。剪枝发生在训练之前。反直觉的一件事是，实际上使用较少的参数从头开始训练模型无法确保能得到一样的结果。尽管一小部分权重比较重要，但你很难提前知道哪些权重是重要的。详情参见论文：《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》。

Ikuyas：

神经网络中的参数缺乏有意义的解释，这是第一堂机器学习课程中就学过的。这并不反直觉，而是预料之中。剪枝后的参数甚至并不被认为是吸收噪声的神经元。对于标准回归模型来说，噪声有时似乎像是正态分布的实现。而神经网络拟合并不假设任何此类事情。因此使用较少的参数，你可以用完全不同的模型拟合数据。剪枝技术并没有什么用。