本文介绍Tensorlfow 针对 CPU SSE4.1 SSE4.2 AVX AVX2 FMA 的编译优化,以提升Tensorflow在CPU上的计算速度,实测可以提升两倍以上的速度。
在用 pip 安装tensorflow的CPU版本后,在运行的时候通常会出现如下提示:Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
该提示说明你的CPU支持AVX扩展,但是你安装的TensorFlow版本无法编译使用。而AVX的利用可以大大提升线性代数运算。
由于AVX不是所有CPU都支持(实际上08年以后的intel CPU都支持...),所以通过 pip install 安装的 tensorflow CPU 版本是不支持AVX的。
---------------
所谓AVX,全称为高级矢量扩展(Advanced Vector eXtensions,AVX)是英特尔在2008年3月提出的英特尔和AMD微处理器的x86指令集体系结构的扩展,英特尔首先通过Sandy Bridge处理器在2011年第一季度推出,随后由AMD推出Bulldozer处理器在2011年第三季度.AVX提供了新功能,新指令和新编码方案。
特别是,AVX引入了融合乘法累加(FMA)操作,加速了线性代数计算,即点积,矩阵乘法,卷积等。几乎所有机器学习训练都涉及大量这些操作,因此将会支持AVX和FMA的CPU(最高达300%)更快。该警告指出您的CPU确实支持AVX。
---------------
Tensorflow使用谷歌开源的 Bazel 自动化构建工具编译项目,有一定的学习成本。
但是可以借助github上编译好的版本进行安装。
(1)linux/Mac OS 编译
github地址: https://github.com/lakshayg/tensorflow-build
下表为开源项目中的一部分,根据Tensorflow、Ubuntu、GCC、和python版本选择对应的下载链接。
实际上版本并没有表中那么严格的限制,比如tensorflow版本,版本 1.4.0~1.13.0的主要API变动不大,假如表格中没有对应版本的话,可以选最近的版本,实测是可以通过的。
同样的,GCC版本假如高于表中的版本,也选最近的低版本即可
操作系统也不一定是Ubuntu,应该Debian系的都可以。
TF | HW | OS | GCC | Python | Supports | |
1.9.0 | CPU | Ubuntu 16.04 | 5.4 | 3.6.6 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
1.9.0 | CPU | Ubuntu 16.04 | 5.4 | 3.5.2 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
1.9.0 | CPU | Ubuntu 16.04 | 5.4 | 2.7.12 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
1.9.0 | CPU | Ubuntu 18.04 | 7.3 | 3.6.5 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
1.10.0 | CPU | Ubuntu 18.04 | 7.3 | 3.6.5 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
1.10.0 | CPU | Ubuntu 18.04 | 7.3 | 2.7.15rc1 | FMA, AVX, AVX2, SSE4.1, SSE4.2 | Download |
下载后得到 tensorflow-1.9.0-cp36-cp36m-linux_x86_64.whl 类似的文件,然后直接在对应的环境中:
pip install tensorflow-1.9.0-cp36-cp36m-linux_x86_64.whl
或者:
conda install tensorflow-1.9.0-cp36-cp36m-linux_x86_64.whl
然后再运行项目,AVX 相关的警告就消失了,而且模型训练和测试速度应该有大幅提升。
(2)windows 编译
github 地址:https://github.com/fo40225/tensorflow-windows-wheel
方法同上,找到对应的版本 pip install 或者conda install 即可