【DL】浅谈模型Inference优化

2,675 阅读3分钟

最近又被一周一更的flag打脸,一是拉来了外援助阵专栏(以后会越来越多的!),二是自己想探索新的故事线(比如NLP+CV的任务),三是工作太忙(懒)。

最近开始接触BERT Inference的优化,这篇文章主要总结一下知识点

1. 简介

Inference提速的重要性就不必说了,目前主流的优化方式有以下几种:

  1. 分布式计算(CPU):将底层的矩阵运算并行化
  2. 混合精度(GPU):理论上INT8优化可提升4倍、FP16可提升2倍
  3. 重写Kernel(GPU):使用GPU计算时,每次运算(比如TF中的operation)都要经过几个流程:CPU在GPU上分配显存 -> CPU把数据发送给GPU -> CPU启动GPU kernel -> CPU把数据取回 -> CPU释放GPU显存。复杂的模型往往需要大量的operation,造成多次显存读写。而如果把多个kernel合并重写就可以避免这种额外开销

接下来会介绍一下Pytorch、Tensorflow、C++配合不同设备可用的优化方案,各位可以根据业务需要和经济实力选择适合自己团队的。

另外在开始正文前先介绍一些周边名词,方便之后的理解:

  • BLAS:Basic Linear Algebra Subprograms,一个线性代数计算的API标准,Netlib用fortran语言编写了一个BLAS计算库
  • LAPACK:基于BLAS,Netlib用fortran语言编写的计算库,优于BLAS
  • Intel MKL:基于BLAS、LAPACK和另外一些计算库编写的数学计算库,C++和Fortran语言编写,基于OpenMP可进行多线程优化。后来针对深度神经网络又出了一个MKL-DNN。Pytorch就是基于MKL的
  • Eigen:基于BLAS、LAPACK和另外一些计算库编写的数学计算库,C++语言编写,基于OpenMP可进行多线程优化。支持MKL作为底层,支持CUDA。MKL和Eigen是两个常用的高级计算库,其中MKL的速度更快一些,但Eigen更容易上手,而且随着优化速度也在逼近MKL。Tensorflow就是基于Eigen的
  • CuDNN:用于深度神经网络的GPU加速库,支持并行计算,Pytorch的GPU实现基于CuDNN,TF也有提供基于cudnn的实现
  • CuBLAS:用cuda实现的GPU BLAS计算库

2. 优化方案(以BERT为例)

2.1 经济CPU

  1. Tensorflow提供C++,Java API,建议把训练好的模型用更高效的语言部署。因为是基于eigen,应该可以设置线程数并行化矩阵运算。
  2. Pytorch好像只有C++ API,通过torchlib调用,可以设置线程数并行化矩阵运算。
  3. 厉害的大佬可以基于MKL或者Eigen自己撸一份,比如知乎的几个大牛就写了一个CuBERT,支持load pb格式的BERT,看性能比Tensorflow提升不少。

2.2 土豪GPU

  1. Tensorflow XLA,可以在编译阶段合并OpKernel
  2. Tensorflow、Pytorch都支持FP16优化
  3. TensorRT是NVDIA提供的库,也是对OpKernel进行融合,同时支持INT8和FP16优化
  4. Faster Transformer是NVIDIA专门针对transformer写的Kernel,极速爆表,支持FP16和TensorRT,提供C++、TF接口,如果拥有GPU的话你还在等什么:

3. 建议

对线上要求很高的话,还是用GPU+faster transformer,如果不高,那也不用自己手撸一份纯C++的代码,用TF、Pytorch的C++ API就可以了