使用 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;
  }
}