TensorFlow 编译优化之 CPU 篇

技术

picture.image


假期一晃而过,估计不少童鞋和小编一样得了假期综合症,明天又要起早贪黑,心塞……

不知有没有童鞋注意到,运行 TensorFlow 时经常会冒出这些蛋疼的警告:

picture.image

抄录如下:

"The TensorFlow library wasn't compiled to use AVX2 instructions, but these are available on your machine and could speed up CPU computations."

翻译一下,就是说 TensorFlow 库编译时没有使用 AVX2 指令集,但你机器上的 CPU 本身是支持 AVX2 的,如果能使用 AVX2 加速,计算效率可以大幅提升。

什么是 AVX2 ? 它有什么用?

简单来说,AVX2, AVX, SSE 都是 Intel CPU 上的 SIMD 指令集,可以一条指令同时处理多个输入。AVX/AVX2 输入位宽可达 256 bits,可以同时处理 8 个 float 或 int,也可以同时处理 16 个 short 或 32 个 char,实现并行加速。对底层指令集感兴趣的童鞋可以阅读【1】深入了解,我们这里不需要手写这些指令,只是加编译选项优化,让编译器自动植入这些指令。

不同 CPU 硬件版本对这些指令集的支持程度有所不同,有些旧的 CPU 可能只支持 AVX,而新的 CPU 可能支持上述全部指令集。Linux 下查看你的 CPU 支持哪些指令集,只需运行 lscpu 就可以看到:

picture.image

上图中 "Intel Core i7-5930K" 就是 CPU 型号,下面的 Flags 有非常多选项,可以找到诸如 mmx, sse, sse3, sse4_1, sse4_2, avx, avx2 ,说明 CPU 支持这些指令集。

回到最开始的问题, 如何让 TensorFlow 能利用这些 SIMD 指令集优化 CPU 上的计算呢? 在 StackOverflow 【2】上找到了解决方法。其实很简单,就是从源码编译 TensorFlow 时,使用如下命令:

$ bazel build -c opt --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-mfpmath=both //tensorflow/tools/pip_package:build_pip_package

其中 "--copt=-mavx" 表示使用 AVX 指令,"--copt=-mavx2" 表示使用 AVX2 指令,后面类似。

完整编译过程可以参考之前的文章《TensorFlow 1.0.0rc1 入坑记》或官网步骤。

在笔者电脑上实测编译优化过的 CPU 版本 TensorFlow 运行 AlexNet Benchmark【3】,Forward + Backward 计算可以提速将近 2 倍!

【1】Intel® Advanced Vector Extensions Programming Reference

【2】 https://stackoverflow.com/questions/41293077/how-to-compile-tensorflow-with-sse4-2-and-avx-instructions

【3】 https://github.com/tensorflow/models/blob/master/tutorials/image/alexnet/alexnet\_benchmark.py


picture.image

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
火山引擎边缘渲染的探索与实践 | 第 11 期边缘云主题Meetup
《火山引擎边缘渲染的探索与实践》黄旭能|火山引擎边缘渲染产品经理
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论