当前位置：首页 > news >正文

从零开始学习tensort模型部署（二）：从文件加载引擎的完整指南

news 2025/11/15 7:09:42

文章目录

前言
一、为什么需要加载TensorRT引擎文件？
二、引擎文件加载过程
- 1.打开engine文件
- 2.获取engine文件大小
- 3.为文件分配内存空间
完整代码示例：

前言

在深度学习模型部署过程中，TensorRT作为NVIDIA推出的高性能推理框架，能够显著提升模型在NVIDIA GPU上的推理速度。今天我们将深入探讨TensorRT模型推理中的一个关键步骤：如何从文件加载预编译的TensorRT引擎。

一、为什么需要加载TensorRT引擎文件？

在正式进入代码解析前，我们先了解一下背景。TensorRT引擎文件（.engine）是通过TensorRT对原始模型（如ONNX、TensorFlow等）进行优化和序列化后生成的二进制文件。这个文件包含了针对特定GPU架构的高度优化代码，直接加载它可以避免每次运行时重新进行优化，大大缩短启动时间。

二、引擎文件加载过程

1.打开engine文件

代码如下（示例）：

std::ifstream file(engine_file_path, std::ios::binary);
assert(file.good());  // 确保文件成功打开

关键点解析：

std::ios::binary标志至关重要，它确保文件以二进制模式打开，避免Windows平台上的换行符转换问题
assert(file.good())提供基本的错误检查，在生产环境中建议使用更健壮的错误处理机制

2.获取engine文件大小

代码如下（示例）：

// 获取文件大小
file.seekg(0, std::ios::end);    // 移动到文件末尾
auto size = file.tellg();         // 获取文件大小
file.seekg(0, std::ios::beg);     // 回到文件开头

注：

这是C++中获取文件大小的标准方法：先定位到文件末尾，获取位置（即文件大小），再回到开头
tellg()返回当前流位置，即文件大小（字节数）

3.为文件分配内存空间

char* trtModelStream = new char[size];
assert(trtModelStream);          // 确保内存分配成功
file.read(trtModelStream, size); // 读取整个文件到内存
file.close();                    // 关闭文件

注：

使用new[]动态分配内存，记得后续要用delete[]释放
file.read()一次性读取整个文件，效率最高
及时关闭文件释放系统资源

完整代码示例：

#include <fstream>
#include <iostream>
#include <cassert>bool loadTRTEngine(const std::string& engine_file_path, char*& trt_model_stream, size_t& engine_size) {std::ifstream file(engine_file_path, std::ios::binary);if (!file.good()) {std::cerr << "无法打开引擎文件: " << engine_file_path << std::endl;return false;}// 获取文件大小file.seekg(0, std::ios::end);engine_size = file.tellg();file.seekg(0, std::ios::beg);if (engine_size == 0) {std::cerr << "引擎文件为空" << std::endl;file.close();return false;}// 分配内存trt_model_stream = new char[engine_size];if (!trt_model_stream) {std::cerr << "内存分配失败" << std::endl;file.close();return false;}// 读取文件内容file.read(trt_model_stream, engine_size);// 检查是否读取了完整文件if (file.gcount() != engine_size) {std::cerr << "文件读取不完整" << std::endl;delete[] trt_model_stream;file.close();return false;}file.close();return true;
}// 使用示例
int main() {const std::string engine_path = "model.engine";char* trt_model_stream = nullptr;size_t engine_size = 0;if (loadTRTEngine(engine_path, trt_model_stream, engine_size)) {std::cout << "成功加载TensorRT引擎，大小: " << engine_size << " 字节" << std::endl;// 这里可以继续创建IRuntime和反序列化引擎// nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);// nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(trt_model_stream, engine_size, nullptr);// 使用完毕后记得清理内存delete[] trt_model_stream;} else {std::cerr << "加载引擎失败" << std::endl;return -1;}return 0;
}