大小仅17KB!小型风格迁移网络包含11686个训练权重

大小仅17KB!小型风格迁移网络包含11686个训练权重现在有很多现成的训练艺术风格迁移模型的工具,大多数人使用Johnson等人描述的网络架构的变体来执行快速的前馈风格化。因此,大多数风格迁移模型都是7MB。

研究表明,神经网络通常比它们需要的大得多,它们所包含的数百万个权重中的许多都是无关紧要的。所以研究者创造仅有11,686个训练权重的17KB神经网络。

大小仅17KB!小型风格迁移网络包含11686个训练权重

左:原始图像。中间:来自小型17KB模型的风格化图像。右:来自较大的7MB模型的风格化图像。

快速概览:

原型:
大小:7MB
权重:iPhone X上的速度:18 FPS

小模型:
大小:17KB
权重:iPhone X上的速度:29 FPS

如何缩小风格迁移模型

事实证明,制作一个小型模型实际上非常简单。研究者最终依赖于两种技术,这两种技术都推广到其他模型。

1.大量修剪层和权重。

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

修剪策略

卷积神经网络通常包含在训练期间调整的数百万甚至数亿个权重。作为一般的经验法则,更多的权重意味着更高的准确性。但是交换效率非常低。谷歌MobileNetV2的库存配置有347万个权重,占用16MB的空间。该InceptionV3架构是2400万点的权重占用92MB大近6倍。尽管包含超过2000万个额外权重,但ImageNet上的InceptionV3排名前1的分类精度仅比MobileNetV2高7个百分点(80%VS 73%)。

因此,如果我们有一个神经网络,我们可以假设大多数权重都不是那么有用并删除它们。但是怎么做?有三种选择:修剪单个权重级别,层级别和块级别。

  • 权重级别:正如我们所见,在某些神经网络中绝大多数(> 95%)训练过的权重都没有帮助。如果我们能够确定哪些权重实际上有助于网络准确性,我们可以保留这些权重并删除其余权重。
  • 层级别:权重打包在单个层中。例如,2D卷积层具有称为内核的权重张量,具有用户定义的宽度,高度和深度。使内核更小会缩小整个网络的大小。
  • 块级别:层通常组合成块,即可重复利用的子图。例如,ResNets名字来源于重复10到50次的“残余块”。块级别的修剪会在一次切割中删除多个层,从而删除参数。

在实践中,稀疏张量操作没有很好的实现,无法使权重级别有价值。希望将来在这方面做得更多。

在实践中修剪

研究者的图层修剪技术是引入宽度乘数作为超参数。谷歌首次在其着名的MobileNet论文中介绍,它既简单又有效。

宽度乘数通过恒定分数调整每个卷积层中的滤波器数量。对于给定的图层和宽度乘数alpha,过滤器的数量F变为alpha * F。

使用这个超参数,可以生成具有相同架构但权重数量不同的连续网络。训练每个配置,我们可以绘制模型的速度和大小与它的准确性之间的权衡。

让我们来看一个构建类似于Johnson等人描述的快速样式传递模型的方法,但这次,宽度乘数被添加为超参数:

@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 ,生成的网络包含1.7M个权重。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个权重的神经网络。任何低于10,000权重的东西都不能持续训练并产生不良的风格化图像。

值得一提的是,这种修剪技术是在网络训练之前应用的。通过在训练期间和训练后进行迭代修剪,您可以在许多任务上获得更好的性能。

量化

最后一段压缩是在网络训练完成之后。神经网络权重通常存储为64或32位浮点数。量化过程将这些浮点权重中的每一个映射到具有较低位宽的整数。从32位浮点权重到8位整数可以将存储大小减少4倍。

现在,每个主要的移动框架都支持量化,包括TensorFlow Mobile,TensorFlow Lite,Core ML和Caffe2Go。

最终结果

总而言之,我们的微型网络架构有11,868个参数,而Johnson的原始模型则为170万个。当转换为Core ML并进行量化时,最终大小仅为17KB,而原始大小为1.7MB,仅为原来的0.10%。以下是梵高星夜的训练结果。

大小仅17KB!小型风格迁移网络包含11686个训练权重

尽管尺寸有400倍的差异,但在 iPhone X 上,小型模型的运行速度仅快了 50%。可能计算与这一通用架构相关,也可能是将图像在GPU上进行处理时造成的。

结论

研究者使用两种简单的技术将风格迁移神经网络的大小减少了99.9%。用简单的宽度乘数超参数修剪层,并且训练的权重从32位浮点数量化到8位整数。将来,这些方法可能会推广到其他神经网络。风格迁移很容易,因为准确性明显可见。对于图像识别这样的更可量化的任务,极端修剪后性能可能明显下降。

代码:github.com/fritzlabs/fritz-style-transfer/blob/master/example/starry_night_640x480_small_a03_q8.mlmodel

本文为ATYUN(www.atyun.com)编译作品,ATYUN专注人工智能
请扫码或微信搜索ATYUN订阅号及时获取最新内容

发表评论