使用 TensorRT 加速带plugin的caffe模型 -- 框架描述

最新推荐文章于 2022-03-05 16:31:10 发布

原创

最新推荐文章于 2022-03-05 16:31:10 发布 · 724 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#TensorRT #plugin #caffe

caffe 模型中含有 TensorRT不支持的层，不需要在prototxt中更改层类型。

以下描述以流程和调用关系为主，具体参数从名称可判断。

流程：

//1、序列化caffe模型。
SerializeCaffeModel();

//2、初始化推断。
InitInference();

for(int i=0; i<N; ++i)
  //3、执行推断。
  DoInference();

//4、释放资源。
FreeInference();

1、序列化caffe模型：

shutdownProtobufLibrary 的调用会导致第二次parse的错误。

a/b/c表示所有plugin层完成a，再执行b、c。

a1/a2/a3表示一个plugin层执行完a1、a2、a3，再下个plugin层执行。

static Logger gLogger;

void SerializeCaffeModel() {
  
  //0、设置gpu设备，后续有cudaMalloc，需要关联设备。
  cudaSetDevice(gpu_id);
  
  //1、创建builder/network/parser。
  IBuilder* builder = createInferBuilder(gLogger);
  INetworkDefinition* network = builder->createNetwork();
  ICaffeParser* parser = createCaffeParser();

  //2、设置插件。
  TrtPluginFa