tensorflow 模型压缩

模型压缩

为了将tensorflow深度学习模型部署到移动/嵌入式设备上,我们应该致力于减少模型的内存占用,缩短推断时间,减少耗电。有几种方法可以实现这些要求,如量化、权重剪枝或将大模型提炼成小模型。

在这个项目中,我使用了 TensorFlow 中的量化工具来进行模型压缩。目前我只使用权重量化来减小模型大小,因为根据 Mac 上的测试结果,完整 8 位转换没有提供额外的好处,比如缩短推断时间。(由于 requant_range 中的错误,无法在 Pixel 上运行完整的 8 位模型)。由于 8 位量化工具不适合 CPU,时间甚至翻了一倍。如果你有兴趣了解更多关于量化的实用建议,可以阅读 Pete Warden 这篇很棒的文章(https://petewarden.com/2017/06/22/what-ive-learned-about-neural-network-quantization/)。

对模型进行权重量化:

  1. 将模型写入协议缓冲区文件。
  2. 从源安装和配置 TensorFlow(https://www.tensorflow.org/install/install_sources)。
  3. 在 TensorFlow 目录下运行下列命令行:
  1. bazel build tensorflow/tools/graph_transforms:transform_graph
  2. bazel-bin/tensorflow/tools/graph_transforms/transform_graph --in_graph=/your/.pb/file --outputs="output_node_name" --out_graph=/the/quantized/.pb/file --transforms='quantize_weights'

以我的项目为例,在量化权重后,预训练的 WaveNet 模型的大小从 15.5Mb 下降到了 4.0Mb。现在可以将这个模型文件移动到安卓项目中的「assets」文件夹。