神经网络推理性能对比:tensorflow, tensorflow-lite, Eigen

cyoahs

背景

笔者正在做一个c++/python混合的项目,其中有一部分涉及升级网络推理。原有网络已经使用tensorflow/keras训练完成。

可选方案

  1. tensorflow-lite:tensorflow-lite模块是tensorflow针对边缘计算和服务部署推出的模块,可以在不同平台和编程语言下进行神经网络的推理,号称做了性能优化,并且可以对模型进一步压缩,提高计算性能。
  2. 原生tensorflow:原生tensorflow在python中使用非常方便,但是在c++中使用时,编译令人秃头。
  3. Eigen:Eigen是C++中最主流的矩阵运算库之一,做了十分极致的优化。

对比

tensorflow vs. tflite

由于作者比较懒,就直接拿实际要用的网络做实验了。网络结构时mix of experts的两级结构,第一部分由两个隐藏层,每层32个神经元,输出N个参数,作为权重。第二部分是专家网络,也是两个隐藏层,但是一共有N个网络,根据前面的N个参数,对参数进行加权平均,再使用平均值后的参数进行计算。专家网络规模较大,论文原文使用了8个专家网络,专家网络每层512个神经元的网络结构。本文使用不同规模的专家网络对tf和tflite的性能进行测试,由于最终应用还是在CPU上跑,所以使用了i5-6500。

(单位:秒)

模型 4-128-128 8-128-128 4-256-256 8-256-256 4-512-512 8-512-512 4-1024-1024 8-1024-1028
tf-lite 0.0035 0.0068 0.0090 0.017 0.030 0.055 0.10 0.19
tf 0.0084 0.0087 0.0090 0.010 0.011 0.013 0.019 0.029

模型一栏的A-B-C表示A个专家,专家网络两个隐藏层的大小分别为B和C。tf使用pip安装,未专门编译,没有AVX加成。
从表格数据可以看出,tf和tf-lite的计算复杂度不同。tf-lite复杂度更高,对小模型性能好,但是scalibility较弱,而tf的scalibility强。猜测tf-lite为了提高性能,阉割了多线程优化部分的功能。

tensorflow vs. Eigen

同样由于作者比较懒,这一部分只做了矩阵乘法的测试,但是结果很显然是有GPU用GPU,没GPU用CPU多线程的tensorflow吊打老老实实单线程跑的Eigen。
实验测试了两个N×N的矩阵乘法性能,时间单位为毫秒。测试平台是双路E5-2696v4/四路GTX 1080Ti,内存管够。tensorflow依然使用pip安装,无AVX加成,C++编译器使用gcc-8。

N 32 64 128 256 512
Eigen -O0 0.287 2.03 10.8 69.5 533
Eigen -O3 3.05 16.6
tf gpu 0.460 0.829
tf cpu 0.573 1.89

不开优化的c++程序运行极慢,使用优化之后性能提高几十倍,测试中O1 O2 O3优化性能都差不多,猜测主要优化了矩阵乘法里面的循环。但是即使如此,还是比tf gpu慢了20倍,而且复杂度更高。同样是cpu,由于tensorflow使用多线程,所以速度仍然比Eigen快,特别是在44C88T这种超多核心的服务器上,对于“小矩阵”的计算甚至能接近GPU。理论上用C++ mkl加上simd优化,应该能接近tf cpu的性能。

小结

综上,原生tensorflow仍然是大型神经网络推理的最佳选择。对于小型神经网络,可以尝试tf lite,特别是内置的模型压缩功能。用C++徒手写矩阵计算的方式吃力不讨好。

阅读 1.6k

1 声望
0 粉丝
0 条评论
你知道吗?

1 声望
0 粉丝
宣传栏