TensorRT深度学习推理框架介绍

一. 产生背景

深度学习的发展带动了一批深度学习框架，caffe、tensorflow、pytorch等，对于计算量庞大的CNN，效率一直是大家所关注的，接触过深度网络压缩的同学应该知道网络压缩最关键的两个思路，剪枝和量化。

TensorRT就是量化，将FP32位权值数据优化为 FP16 或者 INT8，而推理精度不发生明显的降低。

关于TensorRT首先要清楚以下几点：

1. TensorRT是NVIDIA开发的深度学习推理工具，只支持推理，不支持训练；

目前TensorRT3已经支持Caffe、Caffe2、TensorFlow、MxNet、Pytorch等主流深度学习库；

2. TensorRT底层针对NVIDIA显卡做了多方面的优化，不仅仅是量化，可以和 CUDA CODEC SDK 结合使用，

也就是另一个开发包DeepStream；

3. TensorRT独立于深度学习框架，通过解析框架文件来实现，不需要额外安装DL库；

参考示意图：

二. 使用TensorRT

上面是TensorRT的介绍，也可以参考官方文档，更权威一些：https://developer.nvidia.com/tensorrt

下面以Caffe为例介绍TensorRT的使用：

1. caffeToGIEModel －将 caffe model 转换到 TensorRT 格式

＋ void caffeToGIEModel( const std::string& deployFile, // name for caffe prototxt
const std::string& modelFile, // name for model
const std::vector<std::string>& outputs, // network outputs
unsigned int maxBatchSize, // batch size - NB must be at least as large as the batch we want to run with)
IHostMemory *&gieModelStream) // output buffer for the GIE model
{

// 1.创建builder
IBuilder* builder = createInferBuilder(gLogger);

// 2.解析caffe模型，保存到 Network
INetworkDefinition* network = builder->createNetwork();
ICaffeParser* parser = createCaffeParser();
const IBlobNameToTensor* blobNameToTensor = parser->parse(locateFile(deployFile,                 directories).c_str(), locateFile(modelFile, directories).c_str(),*network, DataType::kFLOAT);

// 3.指定输出Tensor
for (auto& s : outputs)
    network->markOutput(*blobNameToTensor->find(s.c_str()));

// 4.构建engine
builder->setMaxBatchSize(maxBatchSize);
builder->setMaxWorkspaceSize(1 << 20);

ICudaEngine* engine = builder->buildCudaEngine(*network);
assert(engine);

// 5.销毁parser
network->destroy();
parser->destroy();

// 6.将engine序列化到GIE，退出
gieModelStream = engine->serialize();
engine->destroy();
builder->destroy();

｝

2. 执行过程 main

// 1.从caffe模型创建GIE模型，序列化到流
IHostMemory *gieModelStream{nullptr};
caffeToGIEModel("mnist.prototxt", "mnist.caffemodel", std::vector < std::string > { OUTPUT_BLOB_NAME }, 1, gieModelStream);

// x.数据获取（略）
// x.解析mean文件(略)

// 2.反序列化，得到Runtime engine 
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(gieModelStream->data(), gieModelStream->size(), nullptr);
if (gieModelStream) gieModelStream->destroy();

// 3.创建上下文
IExecutionContext *context = engine->createExecutionContext();

// 4.运行inference
float prob[OUTPUT_SIZE];
doInference(*context, data, prob, 1);

// 5.销毁engine
context->destroy();
engine->destroy();
runtime->destroy();

3. 推理过程 doInference
void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{

const ICudaEngine& engine = context.getEngine();
// 传递给引擎的输入输出buffer指针- 需要精确的 IEngine::getNbBindings(),这里1个输入＋1个输出
assert(engine.getNbBindings() == 2);
void* buffers[2];

// 1.为了绑定buffer，需要知道输入和输出tensor的names
int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME), 
outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

// 2.创建 GPU buffer 和 stream
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));

// 3.通过DMA 输入到 GPU,  异步之行batch，并通过DMA回传
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE*sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);

// 4.释放 stream 和 buffer
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));

｝

三. 模型转化

TensorRT3.0 虽然号称支持 Caffe、caffe2、TensorFlow、Pytorch等网络模型，实际上例子只提供了Caffe和TensorFlow的直接支持。

对于caffe的支持比较简单，可以直接通过载入deploy file和caffemodel来做，而对于tensorflow则是通过转换为uff格式来加载，可以参考样例程序。

网络模型转换及部署可以分为三个步骤：

1）训练模型并保存为.pb文件；

2）将.pb文件转成.uff格式；

3）利用TensorRT加载并运行模型；

对于caffe2，暂时没找到比较好的转换工具，有经验的朋友可以探讨！

四. 关于统一模型的讨论

在当下深度学习框架肆虐的情况下，tensorflow、pytorch、caffe、caffe2、mxnet、CNTK等等，都有自己的用户群，不同框架之间的格式为交流及数据共享带来了巨大的弊端，各种转换代码，另攻城狮不胜其烦，有竞争必然有结盟，大一统是必然结果。

比如我们来看从 caffe -> caffe2模型格式转换过程

可以采用caffe2提供的转换脚本 caffe_translator.py

python -m caffe2.python.caffe_translator deploy.prototxt pretrained.caffemodel

不同框架之间通过脚本进行格式转换。

NNVM & ONNX

NNVM来自于陈天奇团队，对其比较直观的描述，我们可以参考：

TVM和之前发布的模块化深度学习系统NNVM一起，“组成深度学习到各种硬件的完整优化工具链”。

与LLVM一致，可以理解为将不同语言（DL框架）进行编译、优化、链接的编译器。其目的在于将多个DL框架的用户进行整合，与TF进行PK（TF底层也有同样的功能模块）。

或者说，基于NNVM，你可以实现一套自己的深度学习框架（利用tinyflow），代码只需要2000行，可以实现到各种硬件的快速编译部署，听着就很Tool。

ONNX是Facebook、微软联合推出的一个开放标准，旨在不同框架之间完成互操作；

听起来像NNVM如此宏大，将来是一个什么样的关系，还有待于进一步跟进。

就目前来看，ONNX首先着力于结果模型的转化，也就是推理部分。

（关于标准和入口向来是兵家必争，目前阶段不需要盲目投入过多精力，静等某方胜出方为上册，形势很快明朗）

TensorRT深度学习推理框架介绍

猜你喜欢