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进行加速。