autoware.universe源码略读(3.2)--perception：TensorRT

Overview

依我个人通俗的理解，神经网络的使用可以分为两个部分：使用数据训练和对数据进行测试，那么在实际的应用中，是不应该一边训练一边推理的，或者说不应该从零开始一边训练一边推理的，那么就会使用到预先训练好的模型，TensorRT是可以在NVIDIA各种GPU硬件平台下运行的一个C++推理框架。我们利用Pytorch、TF或者其他框架训练好的模型，可以转化为TensorRT的格式，然后利用TensorRT推理引擎去运行我们这个模型，从而提升这个模型在英伟达GPU上运行的速度。速度提升的比例是比较可观的。（参考：TensorRT详细入门指北，如果你还不了解TensorRT，过来看看吧！）

TensorRT库

比起下载安装包的方法，在安装autoware.universe的依赖库的时候，是有安装了ros-galactic-tensorrt这个库的，安装了这个库应该就可以直接用相关的东西了，这里再看回之前common中的tensorrt_common，这里应该就是为TensorRT库设计的一些接口，这里我们再稍微仔细地看一下。当然这里做的还是准备工作，定义了接口和类，后边可以直接使用

构造函数

构造函数里，涉及到了几个输入的参数

参数名	类型	作用
model_path	string	模型的文件路径
precision	string	模型推理精度。通常是 “fp32”、“fp16” 或 “int8”
calibrator	std::unique_ptr<nvinfer1::IInt8EntropyCalibrator2>	用于 INT8 精度推理的校准器
batch_config	tensorrt_common::BatchConfig	批处理配置
max_workspace_size	size_t	最大工作空间大小。指定 TensorRT 运行时可以使用的最大内存空间，用于优化模型
plugin_paths	std::vector<std::string>	插件路径列表

在构造函数里，核心的代码是下面这两句

  runtime_ = TrtUniquePtr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger_)); // 创建运行对象
  initLibNvInferPlugins(&logger_, ""); // 初始化插件

setup

核心函数，这里执行了主要的模型加载的功能，首先是根据后缀判断，如果是engine就直接调用loadEngine函数来加载这个模型，如果是onnx后缀的话，就先调用buildEngineFromOnnx进行转换，再来加载。

context_ = TrtUniquePtr<nvinfer1::IExecutionContext>(engine_->createExecutionContext());

这句代码是为一个已构建好的推理引擎创建一个执行环境。这个环境用于执行模型推理，包括处理输入数据、运行模型计算以及获取输出结果。所以执行完这句代码后，理论上就可以执行推理的步骤了吧。

loadEngine

这里就是加载模型（引擎可以理解为，有这个东西才能执行后边的步骤，所以也确实算一种引擎吧）的函数。这里的rdbuf()函数是一种缓冲流的操作，返回的是一个指针，把文件加载到缓冲区，然后再给输入输出流之类的，这里是给了字符串流中

std::ifstream engine_file(engine_file_path);  // 打开文件
std::stringstream engine_buffer;  // 创建字符串流
engine_buffer << engine_file.rdbuf();  // 将文件流缓冲区内容插入到字符串流中

下一句代码则是关键，调用的应该就是TensorRT的函数，来创建反序列化引擎的对象

engine_ = TrtUniquePtr<nvinfer1::ICudaEngine>(runtime_->deserializeCudaEngine(
  reinterpret_cast<const void *>(engine_str.data()), engine_str.size()));

buildEngineFromOnnx

这个函数也没什么好说的，就是根据onnx模型构建engine，都是用的库里的函数，具体每一步的注释放在下面的代码里了

bool TrtCommon::buildEngineFromOnnx(
  const std::string & onnx_file_path, const std::string & output_engine_file_path)
{
  // 创建TensorRT构建器
  auto builder = TrtUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger_));
  if (!builder) {
    logger_.log(nvinfer1::ILogger::Severity::kERROR, "Fail to create builder");
    return false;
  }

  const auto explicitBatch =
    1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);

  // 创建网络定义
  auto network =
    TrtUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
  if (!network) {
    logger_.log(nvinfer1::ILogger::Severity::kERROR, "Fail to create network");
    return false;
  }

  // 创建配置
  auto config = TrtUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
  if (!config) {
    logger_.log(nvinfer1::ILogger::Severity::kERROR, "Fail to create builder config");
    return false;
  }

  // 设置配置: 精度和内存池限制
  if (precision_ == "fp16" || precision_ == "int8") {
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
  }
#if (NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 8400
  config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_);
#else
  config->setMaxWorkspaceSize(max_workspace_size_);
#endif

  // 创建ONNX解析器并解析模型文件
  auto parser = TrtUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger_));
  if (!parser->parseFromFile(
        onnx_file_path.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kERROR))) {
    return false;
  }

  // 获取输入维度
  const auto input = network->getInput(0);
  const auto input_dims = input->getDimensions();
  const auto input_channel = input_dims.d[1];
  const auto input_height = input_dims.d[2];
  const auto input_width = input_dims.d[3];

  // 设置优化配置
  auto profile = builder->createOptimizationProfile();
  profile->setDimensions(
    network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMIN,
    nvinfer1::Dims4{batch_config_.at(0), input_channel, input_height, input_width});
  profile->setDimensions(
    network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kOPT,
    nvinfer1::Dims4{batch_config_.at(1), input_channel, input_height, input_width});
  profile->setDimensions(
    network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMAX,
    nvinfer1::Dims4{batch_config_.at(2), input_channel, input_height, input_width});
  config->addOptimizationProfile(profile);

  if (precision_ == "int8" && calibrator_) {
    config->setFlag(nvinfer1::BuilderFlag::kINT8);
    config->setInt8Calibrator(calibrator_.get());
  }

#if TENSORRT_VERSION_MAJOR >= 8
  auto plan =
    TrtUniquePtr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
  if (!plan) {
    logger_.log(nvinfer1::ILogger::Severity::kERROR, "Fail to create host memory");
    return false;
  }
  engine_ = TrtUniquePtr<nvinfer1::ICudaEngine>(
    runtime_->deserializeCudaEngine(plan->data(), plan->size()));
#else
  // 构建TensorRT引擎
  engine_ = TrtUniquePtr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config));
#endif

  if (!engine_) {
    logger_.log(nvinfer1::ILogger::Severity::kERROR, "Fail to create engine");
    return false;
  }

  // save engine
#if TENSORRT_VERSION_MAJOR < 8
  auto data = TrtUniquePtr<nvinfer1::IHostMemory>(engine_->serialize());
#endif
  // 保存引擎到文件
  std::ofstream file;
  file.open(output_engine_file_path, std::ios::binary | std::ios::out);
  if (!file.is_open()) {
    return false;
  }
#if TENSORRT_VERSION_MAJOR < 8
  file.write(reinterpret_cast<const char *>(data->data()), data->size());
#else
  file.write(reinterpret_cast<const char *>(plan->data()), plan->size());
#endif

  file.close();

  return true;
}

权重精度

还是参考TensorRT详细入门指北，如果你还不了解TensorRT，过来看看吧！，前面所提到的精度其实就是神经网络中权重值的数值精度，TensorRT支持FP32、FP16、INT8、TF32等，这几种类型都比较常用。

• FP32：单精度浮点型，没什么好说的，深度学习中最常见的数据格式，训练推理都会用到；
• FP16：半精度浮点型，相比FP32占用内存减少一半，有相应的指令值，速度比FP32要快很多；
• TF32：第三代Tensor Core支持的一种数据类型，是一种截短的 Float32 数据格式，将FP32中23个尾数位截短为10bits，而指数位仍为8bits，总长度为19(=1+8 +10)。保持了与FP16同样的精度(尾数位都是 10 位），同时还保持了FP32的动态范围指数位都是8位)；
• INT8：整型，相比FP16占用内存减小一半，有相应的指令集，模型量化后可以利用INT8进行加速。