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深度学习推理框架介绍_第1张图片

二. 使用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& 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首先着力于结果模型的转化,也就是推理部分。

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

你可能感兴趣的:(深度学习)