用 Intel MKL-DNN 加速 CPU 上的深度学习

Intel最近发布了开源的深度学习软件包 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 = 1USE_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

吐槽和讨论欢迎点这里