用 Intel MKL-DNN 加速 CPU 上的深度学习
· 郑达 Amazon AI Applied ScientistIntel最近发布了开源的深度学习软件包 MKL-DNN,来替换之前的 MKLML。MKL-DNN 专门优化了一系列深度学习里的操作符。现在 Apache MXNet 集成了 MKL-DNN 来加速 CPU 上的深度学习。这次 MXNet 团队和 Intel 团队通过更紧密的合作,不仅提高 MXNet 在 CPU 上的性能,同时也希望让系统比之前使用 MKLML 时更加稳定。鉴于现在 inference 主要还是依赖 CPU,希望这次的优化能给需要 inference 的朋友带来更多的帮助。
MKL-DNN 优化的操作符多用于 CNN 模型,其中包括 Convolution, Inner Product, Pooling, Batch Normalization, Activation。Intel 团队在不久之后会加入 RNN cell 和 LSTM cell 来提升 RNN 模型在 CPU 上的性能。为了得到更好的性能,MKL-DNN 使用了定制的数据格式,这也使得与 MXNet 的集成变得复杂,因为 MXNet 里自带的操作符不能够理解 MKL-DNN 的定制的数据格式。为了不修改 MXNet 里其他的操作符,MXNet 的执行引擎需要能够自动地转变数组里的格式,同时为了得到最好的性能,MXNet 需要在各种操作符混用的情况下尽量减少格式的转换。
现在我将手把手展示如何在 MXNet 里使用 MKL-DNN 来提升性能。推荐的安装方式是直接安装事先编译好的带 MKL-DNN 的 MXNet。
pip install —-pre mxnet-mkl
注意一点是如果之前安装了其他版本,最好先卸载,例如 pip uninstall mxnet
,或者使用虚拟环境安装新版本。
当然用户也可以自己手动编译 MXNet。编译带 MKL-DNN 的 MXNet,用户需要用 安装教程 里的命令来安装 MXNet 需要的软件包。MKL-DNN 编译依赖 cmake,所以用户要额外安装 cmake。编译 MXNet 的时候只需要加上 USE_MKLDNN=1
。
sudo apt-get install -y cmake
make USE_BLAS=openblas USE_MKLDNN=1
安装好带 MKL-DNN 的 MXNet 之后,我们就可以运行 MXNet 上的模型了。因为 MXNet 使用 MKL-DNN 来加速原有的操作符,所以用户并不需要修改任何代码来提升性能。下面我们用 MXNet 自带的 benchmark 来展示使用 MKL-DNN 加速的 MXNet 在 CPU 上的性能。
python example/image-classification/benchmark_score.py
我们这里使用亚马逊云上一台 C5.18xlarge 机器来比较 MXNet 各个版本的性能。
- MXNet-OpenBLAS: 这是 MXNet release 1.1 。这个版本只使用了 OpenBLAS 和 OpenMP 来提速。
- MXNet-MKLML:这是 MXNet-MKL release 1.1 。这个版本使用了 MKLML 来提速。编译这个版本时使用了
USE_MKL2017 = 1
和USE_MKL2017_EXPERIMENTAL = 1
。 - MXNet-MKLDNN:这个版本用 MKL-DNN 来加速,就是前面介绍的的使用
--pre
安装的版本。
为了在多核多处理器机器上有更好的性能,我们需要控制线程的数量,并且绑定线程到 CPU 核上。在 Linux 下,我们可以用下面的环境变量来设置线程。
export KMP_AFFINITY=granularity=fine,compact,1,0
export vCPUs=`cat /proc/cpuinfo | grep processor | wc -l`
export OMP_NUM_THREADS=$((vCPUs / 2))
这里 vCPUs
是亚马逊云上的机器虚拟 CPU 的数量。在 C5.18xlarge 上 vCPUs
是 72。
下面的表格显示了各个版本的 MXNet 在不同的模型上使用不同的批量大小时的性能。这里性能的单位是每秒处理的图片的数量。MXNet-MKLDNN 的性能要比 MXNet 默认的实现快出几十倍,在有些情况下甚至比 MXNet-MKLML 要快出不少。
模型 | 批量大小 | OpenBLAS | MKLML | MKLDNN |
---|---|---|---|---|
VGG-16 | 1 | 1.60 | 75.23 | 80.57 |
8 | 2.52 | 106.24 | 120.19 | |
32 | 2.87 | 121.61 | 106.86 | |
Inception-V3 | 1 | 3.80 | 61.96 | 62.26 |
8 | 4.15 | 154.01 | 162.55 | |
32 | 4.36 | 173.93 | 179.36 | |
Resnet-50 | 1 | 6.40 | 71.70 | 78.98 |
8 | 5.64 | 123.45 | 170.59 | |
32 | 7.28 | 147.24 | 198.56 |