TensorRT 常用函数功能

打印 上一主题 下一主题

主题 1015|帖子 1015|积分 3045

模型推理流程

本节主要看每个函数功能解释,不肯定是最优
1、加载engine模型参数

  1. "step0.通过文件流读取engine文件"
  2. std::vector<char> trtModelStream_;   //engine文件参数读取后保存到该数组中   
  3. size_t size{0};
  4. std::ifstream file(engineFile, std::ios::binary);
  5. if (file.good()) {
  6.     file.seekg(0, file.end);    //设置读取位置为文件的末尾
  7.     size = file.tellg();        //获取到文件的大小, tellg()函数不需要带参数,它返回当前定位指针的位置,也代表着输入流的大小。
  8.     file.seekg(0, file.beg);    //设置读取位置为文件的开头
  9.     trtModelStream_.resize(size);
  10.     file.read(trtModelStream_.data(), size);    //读取文件流的内容保存到数组中
  11.     file.close();
  12. }
  13. "step1.通过RunTime对象反序列化engine引擎"
  14. IRuntime *runtime = createInferRuntime(gLogger);
  15. assert(runtime != nullptr);
  16. m_engine = runtime->deserializeCudaEngine(trtModelStream_.data(), size, nullptr);  //最后一个参数不再使用,必须是一个空指针.
  17. assert(m_engine != nullptr);
  18. "step2.创建context来开辟空间存储中间值,一个engine可以有多个context来并行处理"
  19. //由于引擎保存网络定义和训练参数,因此需要额外的空间。这些保存在context中
  20. m_context = m_engine->createExecutionContext();
  21. assert(m_context != nullptr);
复制代码
2、创建流并预分配GPU缓冲区内存

  1. "步骤一:创建流(可以在推理之前提前创建好)"
  2.     // Pointers to input and output device buffers to pass to engine.
  3.     // Engine requires exactly IEngine::getNbBindings() number of buffers.
  4.     assert(engine.getNbBindings() == 2);    //getNbBindings()获取网络输入输出数量
  5.     void* buffers[2];
  6.     // In order to bind the buffers, we need to know the names of the input and output tensors.
  7.     // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  8.     // 为了绑定缓冲区,我们需要知道输入和输出张量的名称。 请注意,索引必须小于IEngine::getNbBindings()
  9.     "step1.指定输入和输出节点名来获取输入输出索引"
  10.     const int inputIndex = engine.getBindingIndex(m_input_tensor_name.data());//blob名称是转换模型时设置好的,这里需要保持一致
  11.     const int outputIndex = engine.getBindingIndex(m_output_tensor_name.data());
  12.     "step2.在设备上开辟网络输入输出需要的GPU缓冲区(内存)"
  13.     //开辟输入输出需要的GPU内存,由网络输入输出决定
  14.     TENSORRTCHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * m_height * m_width * sizeof(float)));
  15.     TENSORRTCHECK(cudaMalloc(&buffers[outputIndex], batchSize * m_output_size * sizeof(float)));
  16.     "step3.创建流"
  17.     cudaStream_t stream;
  18.     TENSORRTCHECK(cudaStreamCreate(&stream));
复制代码
3、执行模型推理

  1. "调用推理函数(自己创建的)"
  2. float *outdata = new float[m_batch_size * m_output_size];
  3. //注意输入数据传入的方式.ptr的首地址
  4. doInference(*m_context, (float *)pre_img.ptr<float>(0), outdata, m_batch_size);
复制代码
  1. bool TensorRTWrapper::doInference(nvinfer1::IExecutionContext& context, float* input, float* output, int batchSize) {
  2.     const nvinfer1::ICudaEngine& engine = context.getEngine();
  3.    
  4.     "步骤二:执行推理"
  5.     "step1.拷贝数据 从主机(CPU)--->设备(GPU)"
  6.     TENSORRTCHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * m_height * m_width * sizeof(float), cudaMemcpyHostToDevice, stream));
  7.     "step2.执行推理"
  8.     context.enqueue(batchSize, buffers, stream, nullptr);
  9.     "step3.拷贝数据 从设备(GPU)--->主机(CPU)"
  10.     TENSORRTCHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * m_output_size * sizeof(float), cudaMemcpyDeviceToHost, stream));
  11.     "step4.同步流"
  12.     cudaStreamSynchronize(stream);//因为上面的cudaMemcpyAsync()函数是异步方式执行的,所有这里需要进行同步
  13. }
复制代码
4、开释前面创建的对象和内存

  1.     "步骤三:释放内存(可以放在类析构时)"
  2.     cudaStreamDestroy(stream);
  3.     TENSORRTCHECK(cudaFree(buffers[inputIndex]));
  4.     TENSORRTCHECK(cudaFree(buffers[outputIndex]));
复制代码

相关函数功能

bindings

bindings是tensorRT对输入输出张量的形貌,bindings可以认为是个数组,bindings=input-tensor+output-tensor,好比input有a,output有b,c,d。即输入一个tensor,输出3个tensor,那么bindings=[a,b,c,d],bindings[0]=a,bindings[1] =a,bindings[2] =c。此时看到engine->getBindingDimensions(0),0表示索引,你就知道获取的维度是多少了,这里获取的就是a。
GetBindingDimensions

GetBindingDimensions(int index)/GetBindingDimensions(string nodeName) 


  • 获取节点维度接口: 通过端口编号或者端口名称,获取绑定的端口的外形信息.
  • int index: 绑定端口的编号
  • string nodeName: 绑定端口的名称
  • return Dims: 接口返回一个Dims结构体,该结构体包罗了节点的维度大小以及每个维度的详细大小。
getNbBindings()

获取网络 输入+输出 数量,看 bindings 解释
cudaHostAlloc 和 cudaMalloc

CUDA:cudaHostAlloc()-CSDN博客
cudaMalloc / cudaFree 和 cudaMallocAsync / cudaFreeAsync

利用 NVIDIA CUDA 流顺序内存分配器,第 1 部分 - NVIDIA 技能博客
极市开发者平台-计算机视觉算法开发落地平台-极市科技

免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!更多信息从访问主页:qidao123.com:ToB企服之家,中国第一个企服评测及商务社交产业平台。
回复

使用道具 举报

0 个回复

倒序浏览

快速回复

您需要登录后才可以回帖 登录 or 立即注册

本版积分规则

我可以不吃啊

论坛元老
这个人很懒什么都没写!
快速回复 返回顶部 返回列表