jetson orin nano super AI模型部署之路(五)tensorrt C++ api介绍
我们基于tensorrt-cpp-api这个仓库介绍。这个仓库的代码是一个非常不错的tensorrt的cpp api实现,可基于此开发自己的项目。
我们从src/main.cpp开始按顺序说明。
一、首先是声明我们创建tensorrt model的参数。
// Specify our GPU inference configuration options
Options options;
// Specify what precision to use for inference
// FP16 is approximately twice as fast as FP32.
options.precision = Precision::FP16;
// If using INT8 precision, must specify path to directory containing
// calibration data.
options.calibrationDataDirectoryPath = "";
// Specify the batch size to optimize for.
options.optBatchSize = 1;
// Specify the maximum batch size we plan on running.
options.maxBatchSize = 1;
// Specify the directory where you want the model engine model file saved.
options.engineFileDir = ".";
这里Options是在src/engine.h中声明的一个结构体
// 网络的选项配置
struct Options {
// 用于 GPU 推理的精度模式
// 可选值:FP32(单精度浮点)、FP16(半精度浮点)、INT8(8位量化)
Precision precision = Precision::FP16;
// 如果选择 INT8 精度,必须提供校准数据集的目录路径
std::string calibrationDataDirectoryPath;
// 用于计算 INT8 校准数据的批量大小
// 应设置为 GPU 能支持的最大批量大小
int32_t calibrationBatchSize = 128;
// 优化的批量大小
// 表示引擎在构建时会针对该批量大小进行优化
int32_t optBatchSize = 1;
// 最大允许的批量大小
// 表示引擎支持的最大批量大小
int32_t maxBatchSize = 16;
// GPU 设备索引
// 用于指定在哪个 GPU 上运行推理
int deviceIndex = 0;
// 引擎文件保存的目录
// 用于存储 TensorRT 引擎文件
std::string engineFileDir = ".";
// 最大允许的输入宽度
// 默认为 -1,表示期望固定的输入大小
int32_t maxInputWidth = -1;
// 最小允许的输入宽度
// 默认为 -1,表示期望固定的输入大小
int32_t minInputWidth = -1;
// 优化的输入宽度
// 默认为 -1,表示期望固定的输入大小
int32_t optInputWidth = -1;
};
二、构建tensorrt推理的engine
Engine<float> engine(options);
这里Engine是一个模板类,继承字IEngine类。IEngine 是一个抽象接口类,定义了 TensorRT 推理引擎的核心功能。 它包括构建和加载网络、运行推理、获取输入和输出张量维度等功能。 通过模板参数 T,可以支持不同的数据类型(如 float 或 int)。 具体的实现需要在派生类中完成。
IEngine类的实现在include/interfaces/IEngine.h中。
template <typename T>
class IEngine {
public:
// 虚析构函数,确保派生类的资源能够正确释放
virtual ~IEngine() = default;
// 构建并加载 ONNX 模型到 TensorRT 引擎文件
// 参数:
// - onnxModelPath: ONNX 模型文件的路径
// - subVals: 输入数据的减法均值,用于归一化
// - divVals: 输入数据的除法均值,用于归一化
// - normalize: 是否对输入数据进行归一化
// 返回值:
// - 如果成功构建并加载网络,返回 true;否则返回 false
virtual bool buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;
// 从磁盘加载 TensorRT 引擎文件到内存
// 参数:
// - trtModelPath: TensorRT 引擎文件的路径
// - subVals: 输入数据的减法均值,用于归一化
// - divVals: 输入数据的除法均值,用于归一化
// - normalize: 是否对输入数据进行归一化
// 返回值:
// - 如果成功加载网络,返回 true;否则返回 false
virtual bool loadNetwork(std::string trtModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;
// 运行推理
// 参数:
// - inputs: 输入数据,格式为 [input][batch][cv::cuda::GpuMat]
// - featureVectors: 输出数据,格式为 [batch][output][feature_vector]
// 返回值:
// - 如果推理成功,返回 true;否则返回 false
virtual bool runInference(const std::vector<std::vector<cv::cuda::GpuMat>> &inputs,
std::vector<std::vector<std::vector<T>>> &featureVectors) = 0;
// 获取输入张量的维度
// 返回值:
// - 输入张量的维度列表,类型为 nvinfer1::Dims3
virtual const std::vector<nvinfer1::Dims3> &getInputDims() const = 0;
// 获取输出张量的维度
// 返回值:
// - 输出张量的维度列表,类型为 nvinfer1::Dims
virtual const std::vector<nvinfer1::Dims> &getOutputDims() const = 0;
};
再来看Engine类的实现。
template <typename T>
class Engine : public IEngine<T> {
public:
// 构造函数,初始化引擎选项
Engine(const Options &options);
// 析构函数,释放 GPU 缓冲区
~Engine();
// 构建并加载 ONNX 模型到 TensorRT 引擎文件
// 将模型缓存到磁盘以避免重复构建,并加载到内存中
// 默认情况下,输入数据会被归一化到 [0.f, 1.f]
// 如果需要 [-1.f, 1.f] 的归一化,可以通过 subVals 和 divVals 参数设置
bool buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) override;
// 从磁盘加载 TensorRT 引擎文件到内存
// 默认情况下,输入数据会被归一化到 [0.f, 1.f]
// 如果需要 [-1.f, 1.f] 的归一化,可以通过 subVals 和 divVals 参数设置
bool loadNetwork(std::string trtModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) override;
// 运行推理
// 输入格式:[input][batch][cv::cuda::GpuMat]
// 输出格式:[batch][output][feature_vector]
bool runInference(const std::vector<std::vector<cv::cuda::GpuMat>> &inputs, std::vector<std::vector<std::vector<T>>> &featureVectors) override;
// 工具函数:调整图像大小并保持宽高比,通过在右侧或底部填充来实现
// 适用于需要将检测坐标映射回原始图像的场景(例如 YOLO 模型)
static cv::cuda::GpuMat resizeKeepAspectRatioPadRightBottom(const cv::cuda::GpuMat &input, size_t height, size_t width,
const cv::Scalar &bgcolor = cv::Scalar(0, 0, 0));
// 获取输入张量的维度
[[nodiscard]] const std::vector<nvinfer1::Dims3> &getInputDims() const override { return m_inputDims; };
// 获取输出张量的维度
[[nodiscard]] const std::vector<nvinfer1::Dims> &getOutputDims() const override { return m_outputDims; };
// 工具函数:将三维嵌套的输出数组转换为二维数组
// 适用于批量大小为 1 且有多个输出特征向量的情况
static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<std::vector<T>> &output);
// 工具函数:将三维嵌套的输出数组转换为一维数组
// 适用于批量大小为 1 且只有一个输出特征向量的情况
static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<T> &output);
// 工具函数:将输入从 NHWC 格式转换为 NCHW 格式,并应用归一化和均值减法
static cv::cuda::GpuMat blobFromGpuMats(const std::vector<cv::cuda::GpuMat> &batchInput, const std::array<float, 3> &subVals,
const std::array<float, 3> &divVals, bool normalize, bool swapRB = false);
private:
// 构建网络
bool build(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals, bool normalize);
// 将引擎选项序列化为字符串
std::string serializeEngineOptions(const Options &options, const std::string &onnxModelPath);
// 获取设备名称列表
void getDeviceNames(std::vector<std::string> &deviceNames);
// 清除 GPU 缓冲区
void clearGpuBuffers();
// 输入数据的归一化、缩放和均值减法参数
std::array<float, 3> m_subVals{};
std::array<float, 3> m_divVals{};
bool m_normalize;
// 存储输入和输出 GPU 缓冲区的指针
std::vector<void *> m_buffers;
std::vector<uint32_t> m_outputLengths{};
std::vector<nvinfer1::Dims3> m_inputDims;
std::vector<nvinfer1::Dims> m_outputDims;
std::vector<std::string> m_IOTensorNames;
int32_t m_inputBatchSize;
// TensorRT 运行时和推理上下文
std::unique_ptr<nvinfer1::IRuntime> m_runtime = nullptr;
std::unique_ptr<Int8EntropyCalibrator2> m_calibrator = nullptr;
std::unique_ptr<nvinfer1::ICudaEngine> m_engine = nullptr;
std::unique_ptr<nvinfer1::IExecutionContext> m_context = nullptr;
// 引擎选项
const Options m_options;
// TensorRT 日志器
Logger m_logger;
};
// 构造函数:初始化引擎选项
template <typename T>
Engine<T>::Engine(const Options &options) : m_options(options) {}
// 析构函数:清除 GPU 缓冲区
template <typename T>
Engine<T>::~Engine() { clearGpuBuffers(); }
相较于抽象接口类IEngine,Engine类添加了resizeKeepAspectRatioPadRightBottom、transformOutput、blobFromGpuMats这几个static的public函数。
// 工具函数:调整图像大小并保持宽高比,通过在右侧或底部填充来实现
// 适用于需要将检测坐标映射回原始图像的场景(例如 YOLO 模型)
static cv::cuda::GpuMat resizeKeepAspectRatioPadRightBottom(const cv::cuda::GpuMat &input, size_t height, size_t width,
const cv::Scalar &bgcolor = cv::Scalar(0, 0, 0));
// 工具函数:将三维嵌套的输出数组转换为二维数组
// 适用于批量大小为 1 且有多个输出特征向量的情况
static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<std::vector<T>> &output);
// 工具函数:将三维嵌套的输出数组转换为一维数组
// 适用于批量大小为 1 且只有一个输出特征向量的情况
static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<T> &output);
// 工具函数:将输入从 NHWC 格式转换为 NCHW 格式,并应用归一化和均值减法
static cv::cuda::GpuMat blobFromGpuMats(const std::vector<cv::cuda::GpuMat> &batchInput, const std::array<float, 3> &subVals,
const std::array<float, 3> &divVals, bool normalize, bool swapRB = false);
在类中,static 修饰的成员函数是 静态成员函数,它与普通成员函数有以下区别:
(1)静态成员函数不依赖于类的实例 静态成员函数属于类本身,而不是类的某个实例。 调用静态成员函数时,不需要创建类的对象,可以直接通过类名调用。
(2)静态成员函数不能访问非静态成员变量 静态成员函数没有 this 指针,因此无法访问类的非静态成员变量或非静态成员函数。 静态成员函数只能访问类的静态成员变量或调用其他静态成员函数。
在这里为什么使用 static?
(1)工具函数的设计 transformOutput 是一个工具函数,用于将三维嵌套的输出数组转换为二维数组。 它的功能与类的实例无关,因此设计为静态函数更合理。 这样可以直接通过类名调用,而不需要创建类的对象。
(2)提高代码的可读性和效率 将与类实例无关的函数声明为静态函数,可以明确表示该函数不依赖于类的状态。 静态函数的调用效率通常比非静态函数更高,因为它不需要隐式传递 this 指针。
第二个修改是对于getInputDims和getOutputDims函数的处理。在抽象接口类IEngine中,其实现方式是
virtual const std::vector<nvinfer1::Dims3> &getInputDims() const = 0;
在Engine类中,其声明变成了
// 获取输入张量的维度
[[nodiscard]] const std::vector<nvinfer1::Dims3> &getInputDims() const override { return m_inputDims; };
[[nodiscard]] 是 C++17 引入的一个属性,用于提示调用者不要忽略函数的返回值。 如果调用者忽略了带有 [[nodiscard]] 属性的函数的返回值,编译器会发出警告。
在这段代码中,getInputDims() 返回的是输入张量的维度信息。如果调用者忽略了这个返回值,可能会导致程序逻辑错误。 使用 [[nodiscard]] 可以提醒开发者注意返回值的重要性。
例如engine.getInputDims(); // 如果没有使用返回值,编译器会发出警告。
override 是 C++11 引入的关键字,用于显式声明一个函数是从基类继承并重写的虚函数。 如果函数没有正确地重写基类中的虚函数(例如函数签名不匹配),编译器会报错。
它可以防止由于函数签名错误而导致的意外行为。 在这段代码中,getInputDims() 是从基类 IEngine 中继承并重写的虚函数。
第一个 const: 表示返回的引用是常量,调用者不能修改返回的 std::vector。 第二个 const: 表示该成员函数不会修改类的成员变量,保证函数是只读的。
需要注意的是,在src/engine.h中,使用了模版类的实现,这有点要注意:
////////////////////////////// 忽略之前的所有代码
// 构造函数:初始化引擎选项
template <typename T>
Engine<T>::Engine(const Options &options) : m_options(options) {}
// 析构函数:清除 GPU 缓冲区
template <typename T>
Engine<T>::~Engine() { clearGpuBuffers(); }
// Include inline implementations
#include "engine/EngineRunInference.inl"
#include "engine/EngineUtilities.inl"
#include "engine/EngineBuildLoadNetwork.inl"
在文件的最后,才include了这几个头文件。
EngineRunInference.inl: 包含 runInference 函数的实现。
EngineUtilities.inl: 包含工具函数(如 transformOutput)的实现。
EngineBuildLoadNetwork.inl: 包含 buildLoadNetwork 和 loadNetwork 函数的实现。
将 #include 头文件放在文件末尾的原因是为了包含 内联实现文件(.inl 文件),这种设计通常用于模板类或需要将实现与声明分离的情况下。
- .inl 文件的作用 .inl 文件通常用于存放模板类或函数的实现。 在 C++ 中,模板类或模板函数的实现必须在编译时可见,因此不能像普通类那样将实现放在 .cpp 文件中。 为了保持代码的清晰和模块化,开发者会将模板类的实现从头文件中分离出来,放入 .inl 文件中。
- 为什么将 .inl 文件放在文件末尾? 2.1 避免重复定义 如果在头文件的开头包含 .inl 文件,可能会导致重复定义的问题,尤其是在头文件被多次包含时。 将 .inl 文件放在头文件末尾,可以确保模板类的声明先被处理,然后再处理实现部分。
2.2 确保依赖的声明已完成 .inl 文件中的实现可能依赖于头文件中声明的类或函数。 将 .inl 文件放在头文件末尾,可以确保所有必要的声明都已完成,避免编译错误。
2.3 提高代码的可读性 将 .inl 文件放在头文件末尾,可以将类的声明和实现逻辑分开,增强代码的可读性和可维护性。 开发者可以在头文件中快速查看类的接口,而不需要直接阅读实现细节。 3. 为什么不直接放在 .cpp 文件中? 模板类或模板函数的实现不能放在 .cpp 文件中,因为模板的具体类型是在编译时实例化的,而 .cpp 文件是在链接阶段处理的。 如果将模板实现放在 .cpp 文件中,编译器在实例化模板时将无法找到实现,导致链接错误。
三、tensorrt推理Engine的具体实现
3.1 buildLoadNetwork实现
buildLoadNetwork在include/engine/EngineBuildLoadNetwork.inl中实现。
template <typename T>
bool Engine<T>::buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals,
bool normalize) {
const auto engineName = serializeEngineOptions(m_options, onnxModelPath);
const auto engineDir = std::filesystem::path(m_options.engineFileDir);
std::filesystem::path enginePath = engineDir / engineName;
spdlog::info("Searching for engine file with name: {}", enginePath.string());
if (Util::doesFileExist(enginePath)) {
spdlog::info("Engine found, not regenerating...");
} else {
if (!Util::doesFileExist(onnxModelPath)) {
auto msg = "Could not find ONNX model at path: " + onnxModelPath;
spdlog::error(msg);
throw std::runtime_error(msg);
}
spdlog::info("Engine not found, generating. This could take a while...");
if (!std::filesystem::exists(engineDir)) {
std::filesystem::create_directories(engineDir);
spdlog::info("Created directory: {}", engineDir.string());
}
auto ret = build(onnxModelPath, subVals, divVals, normalize);
if (!ret) {
return false;
}
}
return loadNetwork(enginePath, subVals, divVals, normalize);
}
其中,最重要的是build函数的实现。build的实现比较长,该函数的作用是从一个 ONNX 模型文件构建 TensorRT 引擎。 它会解析 ONNX 模型文件,设置优化配置(如动态批量大小、动态输入宽度、精度模式等),并最终生成一个序列化的 TensorRT 引擎文件。
其完整实现如下代码,后面会对其中重要代码展开解释,一并写在代码注释上。
3.2 build函数实现(最重要的函数)
template <typename T>
bool Engine<T>::build(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals, bool normalize) {
// Create our engine builder.
//创建 TensorRT 构建器和网络
// nvinfer1::createInferBuilder:
// 创建一个 TensorRT 构建器(IBuilder),用于构建和优化网络。
// m_logger 是一个自定义的日志器,用于记录 TensorRT 的日志信息。
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(m_logger));
if (!builder) {
return false;
}
// Define an explicit batch size and then create the network (implicit batch
// size is deprecated). More info here:
// https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#explicit-implicit-batch
// builder->createNetworkV2:
// 创建一个网络定义(INetworkDefinition),用于描述神经网络的结构。
// 使用 kEXPLICIT_BATCH 标志表示显式批量大小(TensorRT 7.0 之后推荐使用)。
auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
if (!network) {
return false;
}
// Create a parser for reading the onnx file.
// 创建 ONNX 解析器并解析模型
// nvonnxparser::createParser:
// 创建一个 ONNX 解析器(IParser),用于将 ONNX 模型解析为 TensorRT 的网络定义。
auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, m_logger));
if (!parser) {
return false;
}
// We are going to first read the onnx file into memory, then pass that buffer
// to the parser. Had our onnx model file been encrypted, this approach would
// allow us to first decrypt the buffer.
// 读取 ONNX 文件:
// 将 ONNX 模型文件读入内存缓冲区(buffer),然后传递给解析器。
// 这种方式允许对模型文件进行预处理(如解密)。
std::ifstream file(onnxModelPath, std::ios::binary | std::ios::ate);
std::streamsize size = file.tellg();
file.seekg(0, std::ios::beg);
std::vector<char> buffer(size);
if (!file.read(buffer.data(), size)) {
auto msg = "Error, unable to read engine file";
spdlog::error(msg);
throw std::runtime_error(msg);
}
// parser->parse:
// 解析 ONNX 模型文件,将其转换为 TensorRT 的网络定义。
// ONNX 是一种常见的模型格式,TensorRT 提供了专门的解析器来支持 ONNX 模型。
// 通过将文件读入内存,可以灵活处理加密或压缩的模型文件。
// Parse the buffer we read into memory.
auto parsed = parser->parse(buffer.data(), buffer.size());
if (!parsed) {
return false;
}
// Ensure that all the inputs have the same batch size
const auto numInputs = network->getNbInputs();
if (numInputs < 1) {
auto msg = "Error, model needs at least 1 input!";
spdlog::error(msg);
throw std::runtime_error(msg);
}
// 检查输入的批量大小
// TensorRT 要求所有输入的批量大小必须一致。
const auto input0Batch = network->getInput(0)->getDimensions().d[0];
for (int32_t i = 1; i < numInputs; ++i) {
if (network->getInput(i)->getDimensions().d[0] != input0Batch) {
auto msg = "Error, the model has multiple inputs, each with differing batch sizes!";
spdlog::error(msg);
throw std::runtime_error(msg);
}
}
// Check to see if the model supports dynamic batch size or not
bool doesSupportDynamicBatch = false;
if (input0Batch == -1) {
doesSupportDynamicBatch = true;
spdlog::info("Model supports dynamic batch size");
} else {
spdlog::info("Model only supports fixed batch size of {}", input0Batch);
// If the model supports a fixed batch size, ensure that the maxBatchSize
// and optBatchSize were set correctly.
if (m_options.optBatchSize != input0Batch || m_options.maxBatchSize != input0Batch) {
auto msg = "Error, model only supports a fixed batch size of " + std::to_string(input0Batch) +
". Must set Options.optBatchSize and Options.maxBatchSize to 1";
spdlog::error(msg);
throw std::runtime_error(msg);
}
}
const auto input3Batch = network->getInput(0)->getDimensions().d[3];
bool doesSupportDynamicWidth = false;
if (input3Batch == -1) {
doesSupportDynamicWidth = true;
spdlog::info("Model supports dynamic width. Using Options.maxInputWidth, Options.minInputWidth, and Options.optInputWidth to set the input width.");
// Check that the values of maxInputWidth, minInputWidth, and optInputWidth are valid
if (m_options.maxInputWidth < m_options.minInputWidth || m_options.maxInputWidth < m_options.optInputWidth ||
m_options.minInputWidth > m_options.optInputWidth
|| m_options.maxInputWidth < 1 || m_options.minInputWidth < 1 || m_options.optInputWidth < 1) {
auto msg = "Error, invalid values for Options.maxInputWidth, Options.minInputWidth, and Options.optInputWidth";
spdlog::error(msg);
throw std::runtime_error(msg);
}
}
//设置优化配置
auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
if (!config) {
return false;
}
//设置优化配置,包括动态批量大小和动态输入宽度。
// 使用 IOptimizationProfile 定义输入的最小、优化和最大维度。
// 为什么要这么写?
// TensorRT 允许为动态输入设置优化配置,以便在推理时根据实际输入调整性能。
// 通过设置优化配置,可以在性能和灵活性之间取得平衡。
// Register a single optimization profile
nvinfer1::IOptimizationProfile *optProfile = builder->createOptimizationProfile();
for (int32_t i = 0; i < numInputs; ++i) {
// Must specify dimensions for all the inputs the model expects.
const auto input = network->getInput(i);
const auto inputName = input->getName();
const auto inputDims = input->getDimensions();
int32_t inputC = inputDims.d[1];
int32_t inputH = inputDims.d[2];
int32_t inputW = inputDims.d[3];
int32_t minInputWidth = std::max(m_options.minInputWidth, inputW);
// Specify the optimization profile`
if (doesSupportDynamicBatch) {
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, inputC, inputH, minInputWidth));
} else {
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMIN,
nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, minInputWidth));
}
if (doesSupportDynamicWidth) {
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kOPT,
nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, m_options.optInputWidth));
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMAX,
nvinfer1::Dims4(m_options.maxBatchSize, inputC, inputH, m_options.maxInputWidth));
} else {
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kOPT,
nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, inputW));
optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMAX,
nvinfer1::Dims4(m_options.maxBatchSize, inputC, inputH, inputW));
}
}
config->addOptimizationProfile(optProfile);
// Set the precision level
const auto engineName = serializeEngineOptions(m_options, onnxModelPath);
if (m_options.precision == Precision::FP16) {
// Ensure the GPU supports FP16 inference
if (!builder->platformHasFastFp16()) {
auto msg = "Error: GPU does not support FP16 precision";
spdlog::error(msg);
throw std::runtime_error(msg);
}
config->setFlag(nvinfer1::BuilderFlag::kFP16);
} else if (m_options.precision == Precision::INT8) {
if (numInputs > 1) {
auto msg = "Error, this implementation currently only supports INT8 "
"quantization for single input models";
spdlog::error(msg);
throw std::runtime_error(msg);
}
// Ensure the GPU supports INT8 Quantization
if (!builder->platformHasFastInt8()) {
auto msg = "Error: GPU does not support INT8 precision";
spdlog::error(msg);
throw std::runtime_error(msg);
}
// Ensure the user has provided path to calibration data directory
if (m_options.calibrationDataDirectoryPath.empty()) {
auto msg = "Error: If INT8 precision is selected, must provide path to "
"calibration data directory to Engine::build method";
throw std::runtime_error(msg);
}
config->setFlag((nvinfer1::BuilderFlag::kINT8));
const auto input = network->getInput(0);
const auto inputName = input->getName();
const auto inputDims = input->getDimensions();
const auto calibrationFileName = engineName + ".calibration";
m_calibrator = std::make_unique<Int8EntropyCalibrator2>(m_options.calibrationBatchSize, inputDims.d[3], inputDims.d[2],
m_options.calibrationDataDirectoryPath, calibrationFileName, inputName,
subVals, divVals, normalize);
config->setInt8Calibrator(m_calibrator.get());
}
// CUDA stream used for profiling by the builder.
cudaStream_t profileStream;
Util::checkCudaErrorCode(cudaStreamCreate(&profileStream));
config->setProfileStream(profileStream);
// Build the engine
// If this call fails, it is suggested to increase the logger verbosity to
// kVERBOSE and try rebuilding the engine. Doing so will provide you with more
// information on why exactly it is failing.
// 构建引擎并保存到磁盘
// 使用 builder->buildSerializedNetwork 构建序列化的 TensorRT 引擎。
// 将引擎保存到磁盘,以便后续加载和使用。
// 为什么要这么写?
// 构建引擎是一个耗时操作,将引擎保存到磁盘可以避免重复构建。
// 序列化的引擎文件可以直接加载到内存中,节省时间。
std::unique_ptr<nvinfer1::IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
if (!plan) {
return false;
}
// Write the engine to disk
const auto enginePath = std::filesystem::path(m_options.engineFileDir) / engineName;
std::ofstream outfile(enginePath, std::ofstream::binary);
outfile.write(reinterpret_cast<const char *>(plan->data()), plan->size());
spdlog::info("Success, saved engine to {}", enginePath.string());
Util::checkCudaErrorCode(cudaStreamDestroy(profileStream));
return true;
}
函数的作用
解析 ONNX 模型文件。
设置优化配置(动态批量大小、动态输入宽度等)。
设置精度模式(FP16 或 INT8)。
构建 TensorRT 引擎并保存到磁盘。
涉及的 TensorRT API
更多推荐


所有评论(0)