笔者正在做一个c++/python混合的项目,其中有一部分涉及升级网络推理。原有网络已经使用tensorflow/keras训练完成。
由于作者比较懒,就直接拿实际要用的网络做实验了。网络结构时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为了提高性能,阉割了多线程优化部分的功能。
同样由于作者比较懒,这一部分只做了矩阵乘法的测试,但是结果很显然是有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++徒手写矩阵计算的方式吃力不讨好。