实践教程 | TensorRT部署深度学习模型

作者 | ltpyuanshuai@知乎
来源 | https://zhuanlan.zhihu.com/p/84125533
编辑 | 极市平台

本文仅作学术分享，版权归原作者所有，如有侵权请联系删除。

导读

对需要部署模型的同志来说，掌握用tensorRT来部署深度学习模型的方法是非常有用的。通过Nvidia推出的tensorRT工具来部署主流框架上训练的模型，以便提高模型推断的速度，占用更少的的设备内存。

1.背景

目前主流的深度学习框架（caffe，mxnet，tensorflow，pytorch等）进行模型推断的速度都并不优秀，在实际工程中用上述的框架进行模型部署往往是比较低效的。而通过Nvidia推出的tensorRT工具来部署主流框架上训练的模型能够极大的提高模型推断的速度，往往相比与原本的框架能够有至少1倍以上的速度提升，同时占用的设备内存也会更加的少。因此对是所有需要部署模型的同志来说，掌握用tensorRT来部署深度学习模型的方法是非常有用的。

2.相关技术

上面的图片取自TensorRT的官网，里面列出了tensorRT使用的一些技术。可以看到比较成熟的深度学习落地技术：模型量化、动态内存优化、层的融合等技术均已经在tensorRT中集成了，这也是它能够极大提高模型推断速度的原因。总体来说tensorRT将训练好的模型通过一系列的优化技术转化为了能够在特定平台（GPU）上以高性能运行的代码，也就是最后图中生成的Inference engine。目前也有一些其他的工具能够实现类似tensorRT的功能，例如TVM，TensorComprehensions也能有效的提高模型在特定平台上的推断速度，但是由于目前企业主流使用的都是Nvidia生产的计算设备，在这些设备上nvidia推出的tensorRT性能相比其他工具会更有优势一些。而且tensorRT依赖的代码库仅仅包括C++和cuda，相对与其他工具要更为精简一些。

3. tensorflow模型tensorRT部署教程

实际工程部署中多采用c++进行部署，因此在本教程中也使用的是tensorRT的C++API，tensorRT版本为5.1.5。具体tensorRT安装可参考教程[深度学习] TensorRT安装，以及官网的安装说明。

模型持久化

部署tensorflow模型的第一步是模型持久化，将模型结构和权重保存到一个.pb文件当中。

pb_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), [v.op.name for v in outputs])
with tf.gfile.FastGFile('./pbmodel_name.pb', mode='wb') as f:f.write(pb_graph.SerializeToString())

具体只需在模型定义和权重读取之后执行以上代码，调用tf.graph_util.convert_variables_to_constants函数将权重转为常量，其中outputs是需要作为输出的tensor的列表，最后用pb_graph.SerializeToString()将graph序列化并写入到pb文件当中，这样就生成了pb模型。

生成uff模型

有了pb模型，需要将其转换为tensorRT可用的uff模型，只需调用uff包自带的convert脚本即可。

python /usr/lib/python2.7/site-packages/uff/bin/convert_to_uff.py   pbmodel_name.pb

如转换成功会输出如下信息，包含图中总结点的个数以及推断出的输入输出节点的信息：

tensorRT c++ API部署模型

使用tensorRT部署生成好的uff模型需要先讲uff中保存的模型权值以及网络结构导入进来，然后执行优化算法生成对应的inference engine。具体代码如下，首先需要定义一个IBuilder* builder，一个用来解析uff文件的parser以及builder创建的network，parser会将uff文件中的模型参数和网络结构解析出来存到network，解析前要预先告诉parser网络输入输出输出的节点。解析后builder就能根据network中定义的网络结构创建engine。在创建engine前会需要指定最大的batchsize大小，之后使用engine时输入的batchsize不能超过这个数值否则就会出错。推断时如果batchsize和设定最大值一样时效率最高。举个例子，如果设定最大batchsize为10，实际推理输入一个batch 10张图的时候平均每张推断时间是4ms的话，输入一个batch少于10张图的时候平均每张图推断时间会高于4ms。

IBuilder* builder = createInferBuilder(gLogger.getTRTLogger());
auto parser = createUffParser();
parser->registerInput(inputtensor_name, Dims3(INPUT_C, INPUT_H, INPUT_W), UffInputOrder::kNCHW);
parser->registerOutput(outputtensor_name);INetworkDefinition* network = builder->createNetwork();if (!parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT)){gLogError << "Failure while parsing UFF file" << std::endl;return nullptr;}  builder->setMaxBatchSize(maxBatchSize);builder->setMaxWorkspaceSize(MAX_WORKSPACE);ICudaEngine* engine = builder->buildCudaEngine(*network);if (!engine){gLogError << "Unable to create engine" << std::endl;return nullptr;}

生成engine之后就可以进行推断了，执行推断时需要有一个上下文执行上下文IExecutionContext* context，可以通过engine->createExecutionContext()获得。执行推断的核心代码是：

context->execute(batchSize, &buffers[0]);

其中buffer是一个void*数组对应的是模型输入输出tensor的设备地址，通过cudaMalloc开辟输入输出所需要的设备空间（显存）将对应指针存到buffer数组中，在执行execute操作前通过cudaMemcpy把输入数据（输入图像）拷贝到对应输入的设备空间，执行execute之后还是通过cudaMemcpy把输出的结果从设备上拷贝出来。

更为详细的例程可以参考TensorRT官方的samples中的sampleUffMNIST代码：https://github.com/NVIDIA/TensorRT/tree/master/samples/opensource/sampleUffMNIST

加速比情况

实际工程中我在Tesla M40上用tensorRT来加速过Resnet-50，Inception-resnet-v2，谷歌图像检索模型Delf（DEep Local Features），加速前后单张图推断用时比较如下图（单位ms）：

4. Caffe模型tensorRT部署教程

相比与tensorflow模型caffe模型的转换更加简单，不需要有tensorflow模型转uff模型这类的操作，tensorRT能够直接解析prototxt和caffemodel文件获取模型的网络结构和权重。具体解析流程和上文描述的一致，不同的是caffe模型的parser不需要预先指定输入层，这是因为prototxt已经进行了输入层的定义，parser能够自动解析出输入，另外caffeparser解析网络后返回一个IBlobNameToTensor *blobNameToTensor记录了网络中tensor和pototxt中名字的对应关系，在解析之后就需要通过这个对应关系，按照输出tensor的名字列表outputs依次找到对应的tensor并通过network->markOutput函数将其标记为输出，之后就可以生成engine了。

IBuilder* builder = createInferBuilder(gLogger);INetworkDefinition* network = builder->createNetwork();ICaffeParser* parser = createCaffeParser();DataType modelDataType = DataType::kFLOAT;const IBlobNameToTensor *blobNameToTensor = parser->parse(deployFile.c_str(),modelFile.c_str(),*network,modelDataType);assert(blobNameToTensor != nullptr);for (auto& s : outputs) network->markOutput(*blobNameToTensor->find(s.c_str()));builder->setMaxBatchSize(maxBatchSize);builder->setMaxWorkspaceSize(1 << 30);engine = builder->buildCudaEngine(*network);

生成engine后执行的方式和上一节描述的一致，详细的例程可以参考SampleMNIST

加速比情况

实际工程中我在Tesla M40上用tensorRT加速过caffe的VGG19，SSD速度变为1.6倍，ResNet50，MobileNetV2加速前后单张图推断用时比较如下图（单位ms）

5.为tensorRT添加自定义层

tensorRT目前只支持一些非常常见的操作，有很多操作它并不支持比如上采样Upsample操作，这时候就需要我们自行将其编写为tensorRT的插件层，从而使得这些不能支持的操作能在tensorRT中使用。以定义Upsample层为例，我们首先要定义一个继承自tensorRT插件基类的Upsample类

class Upsample: public IPluginExt

然后要实现该类的一些必要方法，首先是2个构造函数，一个是传参数构建，另一个是从序列化后的比特流构建。

Upsample(int scale = 2) : mScale(scale) {assert(mScale > 0);}
//定义上采样倍数Upsmaple(const void *data, size_t length) {const char *d = reinterpret_cast<const char *>(data), *a = d;mScale = read<int>(d);mDtype = read<DataType>(d);mCHW = read<DimsCHW>(d);assert(mScale > 0);assert(d == a + length);}
~Upsample(){}

一些定义层输出信息的方法：

int getNbOutputs() const override {return 1;}
//模型的输出个数Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) override {// std::cout << "Get ouputdims!!!" << std::endl;assert(nbInputDims == 1);assert(inputs[0].nbDims == 3);return DimsCHW(inputs[0].d[0], inputs[0].d[1] * mScale, inputs[0].d[2] * mScale);}
//获取模型输出的形状。

根据输入的形状个数以及采用的数据类型检查合法性以及配置层参数的方法：

bool supportsFormat(DataType type, PluginFormat format) const override {return (type == DataType::kFLOAT || type == DataType::kHALF || type == DataType::kINT8)&& format == PluginFormat::kNCHW;}
//检查层是否支持当前的数据类型和格式void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,DataType type, PluginFormat format, int maxBatchSize) override{mDtype = type;mCHW.c() = inputDims[0].d[0];mCHW.h() = inputDims[0].d[1];mCHW.w() = inputDims[0].d[2];}
//配置层的参数

层的序列化方法：

size_t getSerializationSize() override {return sizeof(mScale) + sizeof(mDtype) + sizeof(mCHW);}
//输出序列化层所需的长度void serialize(void *buffer) override {char *d = reinterpret_cast<char *>(buffer), *a = d;write(d, mScale);write(d, mDtype);write(d, mCHW);assert(d == a + getSerializationSize());}
//将层参数序列化为比特流

层的运算方法：

size_t getWorkspaceSize(int maxBatchSize) const override {return 0;}
//层运算需要的临时工作空间大小int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace,cudaStream_t stream) override;
//层执行计算的具体操作

在enqueue中我们调用编写好的cuda kenerl来进行Upsample的计算。

完成了Upsample类的定义，我们就可以直接在网络中添加我们编写的插件了,通过如下语句我们就定义一个上采样2倍的上采样层。addPluginExt的第一个输入是ITensor**类别，这是为了支持多输出的情况，第二个参数就是输入个数，第三个参数就是需要创建的插件类对象。

Upsample up(2)；
auto upsamplelayer=network->addPluginExt(inputtensot,1,up)

6.为CaffeParser添加自定义层支持

对于我们自定义的层如果写到了caffe prototxt中，在部署模型时调用caffeparser来解析就会报错。

还是以Upsample为例，如果在prototxt中有下面这段来添加了一个upsample的层：

layer {name: "upsample0"type: "Upsample"bottom: "ReLU11"top: "Upsample1"
}

这时再调用：

const IBlobNameToTensor *blobNameToTensor =   parser->parse(deployFile.c_str(),modelFile.c_str(),*network,modelDataType);

就会出现错误。

之前我们已经编写了Upsample的插件，怎么让tensorRT的caffe parser识别出prototxt中的upsample层自动构建我们自己编写的插件呢？这时我们就需要定义一个插件工程类继承基类nvinfer1::IPluginFactory, nvcaffeparser1::IPluginFactoryExt。

class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactoryExt

其中必须要的实现的方法有判断一个层是否是plugin的方法，输入的参数就是prototxt中layer的name，通过name来判断一个层是否注册为插件。

bool isPlugin(const char *name) override {return isPluginExt(name);}bool isPluginExt(const char *name) override {char *aa = new char[6];memcpy(aa, name, 5);aa[5] = 0;int res = !strcmp(aa, "upsam");return res;
}
//判断层名字是否是upsample层的名字

根据名字创建插件的方法，有两中方式一个是由权重构建,另一个是由序列化后的比特流创建，对应了插件的两种构造函数，Upsample没有权重，对于其他有权重的插件就能够用传入的weights初始化层。mplugin是一个vector用来存储所有创建的插件层。

IPlugin *createPlugin(const char *layerName, const nvinfer1::Weights *weights, int nbWeights) override {assert(isPlugin(layerName));mPlugin.push_back(std::unique_ptr<Upsample>(new Upsample(2)));return mPlugin[mPlugin.size() - 1].get();}
IPlugin *createPlugin(const char *layerName, const void *serialData, size_t serialLength) override {assert(isPlugin(layerName));return new Upsample(serialData, serialLength);}std::vector <std::unique_ptr<Upsample>> mPlugin;

最后需要定义一个destroy方法来释放所有创建的插件层。

void destroyPlugin() {for (unsigned int i = 0; i < mPlugin.size(); i++) {mPlugin[i].reset();}
}

对于prototxt存在多个多种插件的情况，可以在isPlugin，createPlugin方法中添加新的条件分支，根据层的名字创建对应的插件层。

实现了PluginFactory之后在调用caffeparser的时候需要设置使用它，在调用parser->parser之前加入如下代码。

PluginFactory pluginFactory;
parser->setPluginFactoryExt(&pluginFactory);

就可以设置parser按照pluginFactory里面定义的规则来创建插件层，这样之前出现的不能解析Upsample层的错误就不会再出现了。

官方添加插件层的样例samplePlugin：https://github.com/NVIDIA/TensorRT/tree/master/samples/opensource/samplePlugin)可以作为参考。

7.心得体会（踩坑记录）

转tensorflow模型时，生成pb模型、转换uff模型以及调用uffparser时register Input，output，这三个过程中输入输出节点的名字一定要注意保持一致，否则最终在parser进行解析时会出现错误，找不到输入输出节点。
除了本文中列举的pluginExt，tensorRT中插件基类还有IPlugin，IPluginV2，继承这些基类所需要实现的类方法有细微区别，具体情况可自行查看tensorRT安装文件夹下的include/NvInfer.h文件。同时添加自己写的层到网络时的函数有addPlugin，addPluginExt，addPluginV2这几种和IPlugin，IPluginExt，IPluginV2一一对应，不能够混用，否则有些默认调用的类方法不会调用的，比如用addPlugin添加的PluginExt层是不会调用configureWithFormat方法的，因为IPlugin类没有该方法。同样的在还有caffeparser的setPluginFactory和setPluginFactoryExt也是不能混用的。
运行程序出现cuda failure一般情况下是由于将内存数据拷贝到磁盘时出现了非法内存访问，注意检查buffer开辟的空间大小和拷贝过去数据的大小是否一致.
有一些操作在tensorRT中不支持但是可以通过一些支持的操作进行组合替代，比如可替换为，这样可以省去一些编写自定义层的时间。
tensorflow中的flatten操作默认时keepdims=False的，但是在转化uff文时会默认按照keepdims=True转换，因此在tensorflow中对flatten后的向量进行transpose、expanddims等等操作，在转换到uff后用tensorRT解析时容易出现错误，比如“Order size is not matching the number dimensions of TensorRT” 。最好设置tensorflow的reduce，flatten操作的keepdims=True，保持层的输出始终为4维形式，能够有效避免转到tensorRT时出现各种奇怪的错误。
tensorRT中的slice层存在一定问题，我用network->addSlice给网络添加slice层后，在执行buildengine这一步时就会出错nvinfer1::builder::checkSanity(const nvinfer1::builder::Graph&): Assertion `tensors.size() == g.tensors.size()' failed.，构建网络时最好避开使用slice层，或者自己实现自定层来执行slice操作。
tensorRT 的github中有着部分的开源代码以及丰富的示例代码，多多学习能够帮助更快的掌握tensorRT的使用。

8.参考资料

Nvidia TensorRT Samples：https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html

tensorrt-developer-guide：https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/index.html

TensorRT API Docs：https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html

TensorRT Github：https://github.com/NVIDIA/TensorRT

如果觉得有用，就请分享到朋友圈吧！