9.7 KiB
Intel® MKL-DNN on PaddlePaddle: Design Doc
我们计划将英特尔深度神经网络数学库Intel MKL-DNN (Intel Math Kernel Library for Deep Neural Networks)集成到PaddlePaddle, 充分展现英特尔平台的优势,有效提升PaddlePaddle在英特尔架构上的性能。

Figure 1. PaddlePaddle on IA
近期目标
- 完成常用Layer的MKL-DNN实现。
- 完成常见深度神经网络VGG,GoogLeNet 和 ResNet的MKL-DNN实现。
目前的优化,主要针对PaddlePaddle在重构之前的代码框架以及V1的API。 具体的完成状态可以参见这里。
Contents
Overview
我们会把MKL-DNN会作为第三方库集成进PaddlePaddle,与其他第三方库一样,会在编译PaddlePaddle的时候下载并编译MKL-DNN。
同时,为了进一步提升PaddlePaddle在基本数学运算的计算速度,我们也将MKLML即(MKL small library[1]) 作为另一个第三方库集成进PaddlePaddle,它只会包括生成好的动态库和头文件。 MKLML可以与MKL-DNN共同使用,以此达到最好的性能。

Figure 2. PaddlePaddle with MKL Engines
Actions
添加的相关文件和目录结构如下:
PaddlePaddle/Paddle
├── ...
├── cmake/
│ ├── external/
│ │ ├── ...
│ │ ├── mkldnn.cmake
│ │ └── mklml.cmake
└── paddle/
├── ...
├── math/
│ ├── ...
│ └── MKLDNNMatrix.*
└── gserver/
├── ...
├── layers/
│ ├── ...
│ └── MKLDNN*Layer.*
├── activations/
│ ├── ...
│ └── MKLDNNActivations.*
└── tests/
├── ...
├── MKLDNNTester.*
└── test_MKLDNN.cpp
CMake
在CMakeLists.txt中提供一个与MKL有关的总开关:WITH_MKL,它负责决定编译时是否使用MKLML和MKL-DNN
WITH_MKLML控制是否使用MKLML库。 当打开WITH_MKL时,会自动使用MKLML库作为PaddlePaddle的CBLAS和LAPACK库,同时会开启Intel OpenMP用于提高MKLML的性能。WITH_MKLDNN控制是否使用MKL-DNN。 当开启WITH_MKL时,会自动根据硬件配置[2]选择是否编译MKL-DNN。
Matrix
目前在PaddlePaddle中数据都是以nchw的格式存储,但是在MKL-DNN中的排列方式不止这一种。
所以我们定义了一个MKLDNNMatrix用于管理MKL-DNN数据的不同格式以及相互之间的转换。

Figure 3. MKLDNNMatrix
Layers
所有MKL-DNN的Layers都会继承于MKLDNNLayer,该类继承于PaddlePaddle的基类Layer。
在MKLDNNLayer中会提供一些必要的接口和函数,并且会写好forward和backward的基本逻辑,
子类只需要使用定义好的接口,实现具体的函数功能即可。

Figure 4. MKLDNNLayer
每个MKLDNNlayer都会有inVal_,inGrad_,outVal_和outGrad_的MKLDNNMatrix,
分别代表input value, input gradient,output value和output gradient。
它们会存放MKL-DNN用到的internal memory,同时还会定义以ext开头的MKLDNNMatrix(表示external的memory)。
他们主要是当数据格式与PaddlePaddle默认的nchw格式不匹配时,用于转换内存的工作。
必要的转换函数也会在MKLDNNLayer中提前定义好(具体包括reset input、output的value和grad),
这些函数会根据输入参数重新设置internal和external的memory(当然这两者也可以相等,即表示不需要转换),
每个MKLDNNlayer的子类只需要使用internal的memory就可以了,所有external的转换工作都会在reset函数中都准备好。
一般来说,每个MKLDNNLayer中的extOutVal_和extOutGrad_必须分别与output_.value和output_.grad共享内存,
因为PaddlePaddle的activation会直接使用output_.value和output_.grad,
如果不需要external的buffer用于转换,那么internal的buffer也会与它们共享内存。
Activations
在重构前的PaddlePaddle中,激活函数是独立于Layer的概念,并且输入输出都是公用一块内存,
所以添加了对应的MKLDNNActivation来实现,方式类似于MKLDNNLayer。
Parameters
对于有参数的层,我们会保证MKLDNNLayer使用的参数与PaddlePaddle申请的buffer公用一块内存。
如果存在数据排列格式不一样的情况时,我们会在网络训练之前把格式转换为MKL-DNN希望的格式,
在训练结束的时候再保存为PaddlePaddle的格式,但是整个训练过程中不需要任何转换。
这样既使得最终保存的参数格式与PaddlePaddle一致,又可以避免不必要的转换。
Gradients
由于MKL-DNN的操作都是直接覆盖的形式,也就是说输出的结果不会在原来的数据上累加,
这样带来的好处就是不需要一直清空memory,节省了不必要的操作。
但是注意的是,当网络出现分支且在backward的时候,需要累加不同Layer传过来的梯度。
所以在MKLDNNlayer中实现了一个merge的方法,此时每个小分支的Input Gradient
会先临时保存在MKLDNNMatrix中,由分支处的Layer负责求和,并把结果放到当前层的output_.grad中。
所以整体上,在实现每个子类的时候就不需要关心分支的事情了。

Figure 5. Merge Gradients
Unit Tests
我们会添加test_MKLDNN.cpp和MKLDNNTester.*用于MKL-DNN的测试。
测试分为每个Layer(或Activation)的单元测试和简单网络的整体测试。
每个测试会对比PaddlePaddle中CPU算出的结果与MKL-DNN的结果,小于某个比较小的阈值认为通过。
Protobuf Messages
根据具体layer的需求可能会在proto/ModelConfig.proto里面添加必要的选项。
Python API
目前只考虑v1 API。
计划在python/paddle/trainer/config_parser.py里面添加use_mkldnn这个选择,方便用户选择使用MKL-DNN的layers。
具体实现方式比如:
use_mkldnn = bool(int(g_command_config_args.get("use_mkldnn", 0)))
if use_mkldnn
self.layer_type = mkldnn_*
所有MKL-DNN的layer_type会以*mkldnn_*开头,这些会在MKLDNN*Layer注册layer的时候保证,以示区分。
同时,会在paddle/utils.Flags中添加一个use_mkldnn的flag,用于选择是否使用MKL-DNN的相关功能。
Demos
可能会在v1_api_demo目录下添加一个mkldnn的文件夹,里面放入一些用于MKL-DNN测试的demo脚本。
Benchmarking
会添加benchmark/paddle/image/run_mkldnn.sh,用于测试和对比,在使用MKL-DNN前后的性能。
Others
- 如果在使用MKL-DNN的情况下,会把CPU的Buffer对齐为4096,具体可以参考MKL-DNN中的memory。
- 深入PaddlePaddle,寻找有没有其他可以优化的可能,进一步优化。比如可能会用OpenMP改进SGD的更新性能。
Design Concerns
为了更好的符合PaddlePaddle的代码风格[3],同时又尽可能少的牺牲MKL-DNN的性能[4]。
我们总结出一些特别需要注意的点:
- 使用deviceId_。为了尽可能少的在父类Layer中添加变量或者函数,
我们决定使用已有的
deviceId_变量来区分layer的属性,定义-2为MKLDNNLayer特有的设备ID。 - 重写父类Layer的init函数,修改
deviceId_为-2,代表这个layer是用于跑在MKL-DNN的环境下。 - 创建
MKLDNNBase,定义一些除了layer和memory相关的类和函数。 包括MKL-DNN会用到MKLDNNStream和CPUEngine,和未来可能还会用到FPGAEngine等。 - 如果MKL-DNN layer的后面接有cpu device,那么就会使
output_.value与extOutVal_共享内存, 同时数据格式就是nchw,这样下一个cpu device就能拿到正确的数据。 在有cpu device的时候,external的memory的格式始终是nchw或者nc。
References
- MKL small library是Intel MKL的一个子集。 主要包括了深度学习相关的数学原语与操作,一般由MKL-DNN在发布新版本时一起更新。
- MKL-DNN System Requirements。 目前在PaddlePaddle中,仅会在支持AVX2指令集及以上的机器才使用MKL-DNN。
- 原来的方案会引入nextLayer的信息。 但是在PaddlePaddle中,无论是重构前的layer还是重构后的op,都不会想要知道next layer/op的信息。
- MKL-DNN的高性能格式与PaddlePaddle原有的
NCHW不同(PaddlePaddle中的cuDNN部分使用的也是NCHW,所以不存在这个问题)。 所以需要引入一个转换方法,并且只需要在必要的时候转换这种格式,才能更好的发挥MKL-DNN的性能。