使用 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、设置插件。
  TrtPluginFactory trtpluginfactory;
  parser->setPluginFactoryExt(&trtpluginfactory);

  //3、解析网络和权重。
  const IBlobNameToTensor* blobNameToTensor = parser->parse(net_filename, weights_filename, *network, DataType::kFLOAT);
    //3-a1、调用 nvcaffeparser1::IPluginFactory 的 virtual IPlugin* createPlugin(const char* layerName, const Weights* weights, int nbWeights)
    //3-a2、调用 IPlugin 的 int getNbOutputs() const override
    //3-a3、调用 IPlugin 的 Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override

  //4、标记输出。
  for (int i = 0; i < output_blob_num; ++i) {
    ITensor* tensor = blobNameToTensor->find(output_blob_name[i]);
    network->markOutput(*tensor);
  }

  //5、设置最大batchsize（影响序列化后模型文件的大小）/最大工作大小/数据模式。
  builder->setMaxBatchSize(nMaxBatchSize);
  builder->setMaxWorkspaceSize(nMaxWorkspaceSize);
  builder->setInt8Mode(false);
  builder->setFp16Mode(false);
  
  //6、创建engine。
  ICudaEngine* engine = builder->buildCudaEngine(*network);
    //6-a1、调用 PluginExt 的 bool supportsFormat(DataType type, PluginFormat format) const override
    //6-b1、调用 PluginExt 的 void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) override
    //6-b2、调用 IPlugin 的 size_t getWorkspaceSize(int maxBatchSize) const override
    //6-c1、调用 IPlugin 的 int initialize() override

  //7、序列化。
  IHostMemory* serializedmodel = engine->serialize();
    //7-d1、调用 IPlugin 的 size_t getSerializationSize() override
    //7-d2、调用 IPlugin 的 void serialize(void* buffer) override
  FILE* fp = fopen(sm_filename, "wb");
  if (fp) {
    fwrite((int)serializedmodel->size(), sizeof(int), 1, fp);
    fwrite(serializedmodel->data(), serializedmodel->size(), 1, fp);
    fclose(fp);
  }

  //8、释放对象。
  serializedmodel->destroy();
  network->destroy();
  parser->destroy();
  engine->destroy();
    //8-a、调用 IPlugin 的 void terminate() override
  builder->destroy();
  trtpluginfactory.destroyPlugin();
    //8-b、调用 TrtPluginFactory 的 void destroyPlugin()
  
  //确定不调用parse时，可调用。
  //shutdownProtobufLibrary();
}

2、初始化推断：

定义为全局或成员变量为了后续使用和资源释放。

ICudaEngine* g_engine = nullptr;
IExecutionContext* g_context = nullptr;;
TrtPluginFactory* g_factory = nullptr;;

void InitInference() {
  
  //0、加载序列化模型。
  FILE* fp = fopen(sm_filename, "rb");
  int sm_size;
  float* sm_data = nullptr; 
  if (fp) {
    fread(&sm_size, sizeof(int), 1, fp);
    sm_data = new float[sm_size];
    fread(sm_data, sm_size, 1, fp);
    fclose(fp);
  }
  
  cudaSetDevice(gpu_id);
  
  //1、创建runtime。
  IRuntime* runtime = createInferRuntime(gLogger);
  
  //2、创建engine。
  g_factory = new TrtPluginFactory;
  g_engine = runtime->deserializeCudaEngine((const void*)sm_data, sm_size, (TrtPluginFactory*)m_pvPluginFactory);
    //2-a1、调用 nvinfer1::IPluginFactory 的 IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override
    //2-b1、调用 IPlugin 的 int initialize() override
    
  //3、创建context。
  g_context = g_engine->createExecutionContext();

  //4、销毁runtime。
  runtime->destroy();
  
  delete[] sm_data;
  sm_data = nullptr;
  return;
}

3、执行推断：

void DoInference() {
  
  //1、分配buff。
  const ICudaEngine& engine = g_context->getEngine();
  
  int buff_num = output_blob_num + 1;
  float** ppfBuffers = new float*[buff_num];
 
  cudaMalloc((void**)&ppfBuffers[0], batch_size * input_size);
  for (int i = 0; i < output_blob_num; ++i) {
    const int index = engine.getBindingIndex(output_blob_name[i]);
    cudaMalloc((void**)&ppfBuffers[index], batch_size * output_blob_size[i]);
  }
 
  2、创建stream。
  // Create stream
  cudaStream_t stream;
  cudaStreamCreate(&stream);
  
  3、赋值输入。
  cudaMemcpyAsync(ppfBuffers[0], batch_input_data, batch_size * input_size, cudaMemcpyHostToDevice, stream));
  cudaStreamSynchronize(stream);
   
  4、前向传播。
  g_context->enqueue(batch_size, (void**)ppfBuffers, stream, nullptr);
  
  5、获取输出。
  for (int i = 0; i < output_blob_num; ++i) {
    const int index = engine.getBindingIndex(output_blob_name[i]);
    cudaMemcpyAsync(batch_output_data[i], ppfBuffers[index], batch_size * output_blob_size[i], cudaMemcpyDeviceToHost, stream));
    cudaStreamSynchronize(stream);
  }
  
  6、释放资源。
  cudaStreamDestroy(stream);    
  for (int i = 0; i < buff_num; ++i)
    cudaFree(ppfBuffers[i]);
  delete[] ppfBuffers;
  
  return;
}

4、释放资源：

释放顺序错误会触发异常。

void FreeInference() {
  //释放顺序很重要。
  if (g_context) {
    g_context->destroy();
    g_context = nullptr;
  }

  if (g_engine) {
    g_engine->destroy();
    g_engine = nullptr;    
  }

  if (g_factory) {
    g_factory->destroyPlugin();
    delete g_factory;
    g_factory = nullptr;
  }
}