大小仅17KB!这个微型风格迁移模型太好玩了|代码+教程

301 阅读5分钟
原文链接: zhuanlan.zhihu.com
铜灵 编译整理
量子位 出品 | 公众号 QbitAI

一个风格迁移模型可以有多小?

17KB足矣。

事实证明,你可以用1.7M的参数训练出风格迁移模型,然后只留下11868个参数,生成迁移图像。将模型量化后,原本7M的模型压缩到只有17KB,而迁移效果和原本模型差别不大:

△ 迁移效果对比 | 从左到右分别为原图,小模型效果,大模型效果

17KB是什么概念?一个模糊得不行的小头像,word中的10个文字,迅雷0.01秒内下载的文件……

整波操作效果明显——

原始模型

  • 大小:7M
  • 权重数量:1.7M
  • iPhoneX上的运行速度:18 FPS

缩减后的模型

  • 大小:17KB
  • 权重数量:11868
  • iPhoneX上的运行速度:29 FPS

这到底怎么做到的?

只需两步

这项研究的作者、AI部署解决方案Fritz的联合创始人Jameson Toole在博客中表示,缩小模型真的“pretty easy”,只要你:

一是要毫不留情地修剪层数和权重,二是要通过量化的方式,将32位的浮点权重转换成8位整数。下面分别细讲。

修剪策略

神经网络的训练期间需要调整数百万甚至数亿个权重,一般来讲更多的权重意味着更高的准确性,但非常低效。

假设一个神经网络可以大多数权重都没那么有用,可以删除,应该在什么层面去删?Toole认为主要有三种选择:修剪单个权重、层和模块等。

权重方面:指一些网络中绝大多数(可能会>95%)训练过的权重对结果没有帮助,确定哪些权重能提升网络准确性,将无用权重删去。

层级方面:每个层中都有权重,2D的卷积层具有的权重张量被称为一个卷积核(kernel),包含宽度、高度和深度信息,缩小卷积核也是缩小模型的一种方法。

模块方面:层通常会被组合在一起,形成可重复使用的子图,也就是一个模块。在这个层面的修剪会一次删除多个层,从而减小模型体积。

具体操作

作者引入了宽度乘数(width multiplier)当作超参数,这种方法是谷歌首次在论文MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications中提出的。

论文地址:

https://arxiv.org/abs/1704.04861

宽度乘数能调整每个卷积层中滤波器的数量。如果层和宽度乘数alpha已经给定,那将滤波器的数量F改成F×alpha即可。

通过这个超参数,我们能够利用李飞飞团队曾经发表的方法,生成架构相同但权重数量不同的连续网络,就像下面这样:

李飞飞论文Perceptual Losses for Real-Time Style Transfer and Super-Resolution:

https://arxiv.org/abs/1603.08155

当alpha=1.0时,生成网络的权重为1.7M,当alpha=0.5时,得到权重数量为424102个的一个网络。

此时可以使用低宽度乘数制作构建一些小网络,但里面也有很多重复模块。

作者发现,没有办法大刀阔斧地修剪这些重复模块,参数数量固定时,更深的网络产生的结果更好。作者最终将五个残差块中删去了两个,并将每层的默认滤波器减少到32个。

看起来是这样的:

用上述架构,作者在反复试验后将宽度乘数调整到了0.3,每层中有9个滤波器,最终生成一个只有11868个权重的神经网络。

量化

神经网络训练完成后进入最后的压缩环节。

通常大家会将网络权重存储成64位或32位的浮点数,所谓量化,就是将每一个浮点权重映射成位宽较低的整数。

从32位到8位的映射能将存储空间减少4倍,并且目前每个主流移动端框架都支持量化,包括TensorFlow Mobile、TensorFlow Lite、Core ML和Caffe2Go。

结果展示

最终,用11868个参数构建的这个微型网络架构原始大小为1.7M,当用Apple的移动端框架CoreML量化后,最终大小为17KB,仅为原始大小的0.25%。

这个超小型的迁移模型,可以将所见场景转化成梵高名画《星空》的风格。

作者还发现,尽管模型大小缩小了400倍,但在iPhoneX上的运行速度只提高了50%,原因尚不明确。

用简单两步操作将风格迁移模型大小减少99.75%的方法,这次你学会了吗?

传送门

可进一步阅读博客原文:

https://heartbeat.fritz.ai/creating-a-17kb-style-transfer-model-with-layer-pruning-and-quantization-864d7cc53693

GitHub上可下载模型的代码+Demo:

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

欢迎大家关注我们的专栏:量子位 - 知乎专

诚挚招聘

量子位正在招募编辑/记者,工作地点在北京中关村。期待有才气、有热情的同学加入我们!

相关细节,请在量子位公众号(QbitAI)对话界面,回复“招聘”两个字。

量子位 QbitAI· 头条号签约作者

վ'ᴗ' ի 追踪AI技术和产品新动态