本教程采用的TensorRT版本为10.14.1。

一、导出待部署模型

        为了构建引擎,我们首先需要将我们的模型导出为ONNX格式(其他格式TensorRT目前已经很少维护,逐渐被ONNX取代,如UFF,Caffe等)。主流框架都支持导出模型为ONNX格式,以pytorch为例:

import torchvision.models as models
import torch
import torch.onnx

# load the pretrained model
resnet50 = models.resnet50(pretrained=True, progress=False).eval()

# creat a dummy input batch
BATCH_SIZE=32
dummy_input=torch.randn(BATCH_SIZE, 3, 224, 224)

# export the model to ONNX
torch.onnx.export(resnet50, dummy_input, "resnet50_pytorch.onnx", verbose=False)

二、构建序列化引擎

        通过利用ONNX模型构建引擎,我们就得到了TensorRT优化后的可执行推理计划,该计划是一个二进制文件,通常以.engine、.trt或者.plan作为后缀,其实就是个二进制文件,本质都一样。先来理解一下引擎文件中包含哪些内容,它包含优化后的计算图、固化的权重(量化和适配GPU的存储方式存储)、硬件适配的算子实现、内存布局等,是可以直接在指定 GPU 上运行的 “编译后推理程序”,而不仅仅是重组后的模型和权重。  

        TensorRT 构建引擎主要有三种核心方法,按 “易用性 - 灵活性” 维度划分,覆盖从快速验证到深度定制的全场景:

方法类型 核心原理 实现方式 典型工具 / API
命令行工具:trtexec 基于官方封装的命令行工具,一键式将 ONNX/UFF 转换为引擎 无需写代码,终端执行命令 TensorRT 自带的trtexec.exe(Windows)/trtexec(Linux)
解析器 API  通过 TensorRT 的解析器(ONNX/UFF/CAFFE)直接解析模型文件,在代码中构建引擎 C++/Python 代码调用解析器 + Builder API nvonnxparser(ONNX 解析器)、nvuffparser(UFF 解析器)
网络定义 API 手动 完全不依赖模型文件,手动用 TensorRT API 逐层定义网络结构、设置权重,再构建引擎 C++/Python 代码手动拼接每一层算子

nvinfer1::INetworkDefinition API

方法类型 优点 缺点 适用场景
命令行工具:trtexec

① 零代码,一键生成引擎

② 调试友好,输出详细日志

③ 参数丰富,支持FP16/INT8,动态维度等所有核心配置

① 灵活性差,仅支持离线构建

② 定制化弱,无法自定义算子

① 模型部署前期快速验证和调试

② 离线预生成引擎

解析器 API 

① 能嵌入到业务代码中,支持动态构建或者离线构建+在线加载

② 覆盖主流场景,支持ONNX,适配Pytorch/TensorFlow导出的模型

① 依赖解析器兼容性,若模型含自定义算子需编写解析器插件

② 需要写基础的Builder/Config代码

① 工业级部署方案

② 需要在代码中动态配置引擎参数

③ 需集成自定义算子

网络定义 API 手动

① 极致灵活,完全掌握网络定义

② 无解析器依赖

① 开发成本极高

② 仅适合简单场景

① 资源极度受限

② 需深度定制网络

1、trtexec.exe构建引擎

它利用 TensorRT ONNX 解析器将 ONNX 模型加载到 TensorRT 网络图中,并利用 TensorRT Builder API 生成优化的引擎

trtexec --onnx=fcn-resnet101.onnx --saveEngine=fcn-resnet101.engine --memPoolSize=workspace:8192

参数解析:

(1) 模型输入 / 输出(核心必选区)

参数 说明 是否必选 示例
--onnx=<file> 指定输入的 ONNX 模型文件路径 ✅(转换 ONNX 时必选) --onnx=nnunet_model_with_argmax.onnx
--loadEngine=<file> 加载已有的 TensorRT 引擎文件(无需转换,直接测试) ✅(加载引擎时必选) --loadEngine=nnunet_engine.plan
--saveEngine=<file> 将转换后的 TensorRT 引擎保存到指定路径 ❌(但转换时建议必加) --saveEngine=nnunet_engine.plan
--output=<name> 指定模型输出节点名称(多输出时需指定,单输出可省略) --output=output

(2) 输入形状配置(针对动态形状/固定形状)

参数 说明 是否必选 示例
--shapes=<input_name>:<shape> 推理时指定输入张量的固定形状,构建时固定动态输入的形状(静态形状) --shapes=input:1x1x28x256x256
--minShapes=<input_name>:<shape> 动态形状的最小尺寸(需和 --opt/--max 配合) ✅(动态形状时必选) --minShapes=input:1x1x16x128x128
--optShapes=<input_name>:<shape> 动态形状的最优尺寸(推理时优先用) ✅(动态形状时必选) --optShapes=input:1x1x28x256x256
--maxShapes=<input_name>:<shape> 动态形状的最大尺寸 ✅(动态形状时必选) --maxShapes=input:1x1x32x512x512
--inputIOFormats=<format> 指定输入数据格式(如 FP32:CHW) --inputIOFormats=fp32:chw
--outputIOFormats=<format> 指定输出数据格式 --outputIOFormats

(3) 精度配置

参数 说明 是否必选 示例
--fp32 使用 FP32 精度(默认) --fp32
--fp16 使用 FP16 精度(GPU 需支持,算力大于=5.3,速度提升) --fp16
--int8 使用 INT8 精度(需校准,速度最快) --int8 --calib=<calib_file>
--best 自动选择最优精度(FP16/FP32/INT8) --best

(4)日志/调试参数

参数 说明 是否必选 示例
--verbose 输出详细日志(调试用) --verbose
--exportProfile=<file> 导出性能分析文件(JSON 格式) --exportProfile=profile.json
--dumpOutput 输出推理结果(验证正确性) --dumpOutput
--dumpProfile 打印每层的性能耗时 --dumpProfile

(5)硬件/引擎优化参数

参数 说明 是否必选 示例
--device=<N> 指定使用的 GPU 设备 ID(多卡时) --device=0
--memPoolSize=workspace:<N> 设置工作空间大小(MB),默认 16,越大尝试优化策略越多,只与构建引擎有关。 --workspace=1024/8192
--builderOptimizationLevel=<N> 引擎优化级别(0-5,5 最优,默认2) --builderOptimizationLevel=5

(5)推理性能测试

参数 说明 10.x 优化
--batch=<size> 指定推理批次大小(静态批次) 10.x 支持批次自动适配
--iterations=<num> 指定推理迭代次数(默认 10 次) 建议设为 1000,结果更稳定:--iterations=1000
--duration=<sec> 指定推理持续时间(秒),优先级高于 --iterations --duration=30(测试 30 秒)
--warmUp=<sec> 预热时间(秒),避免初始推理速度波动 --warmUp=5(预热 5 秒)
--avgRuns=<num> 计算平均性能的运行次数 --avgRuns=10(取 10 次平均)
--exportProfile=<path> 导出层级性能分析文件(JSON 格式) 10.x 支持更详细的层耗时统计

总结:

  • 核心必选参数
    • 转换 ONNX 时:--onnx(输入模型) + --shapes(固定形状)/--min/opt/maxShapes(动态形状);
    • 加载引擎时:--loadEngine(输入引擎文件);
    • 保存引擎时(建议必加):--saveEngine
  • 高频可选参数
    • 精度优化:--fp16
    • 性能测试:--iterations--warmUp
    • 内存优化:--workspace(建议设为 1024-4096)。

2、解析器API构建引擎

        利用解析器API构建引擎可以动态的根据当前设备构建引擎,避免引擎与设备不兼容导致的问题。构建引擎可以分为以下几步:

  1. 创建构建器:负责管理引擎构建过程,包括硬件检查、内存分配等。
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(
        nvinfer1::createInferBuilder(*logger));
  2. 创建网络定义(Network):创建一个空的计算图结构,后续通过解析器填充。

    const auto explicitBatch = 1U << static_cast<uint32_t>(
        nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
        builder->createNetworkV2(explicitBatch));
  3. 创建解析器(Parser):解析ONNX模型文件,将ONNX算子映射到TensorRT层并处理常量折叠等优化,实现了ONNX到TensorRT网络的转换。

    auto parser = std::unique_ptr<nvonnxparser::IParser>(
        nvonnxparser::createParser(*network, *logger));
  4. 解析ONNX模型:递归遍历ONNX计算图,为每个节点创建对应的TensorRT层。

    if (!parser->parseFromFile(onnxPath.c_str(),
        static_cast<int>(TrtLogger::Severity::kWARNING))) {}
  5. 创建配置(config):配置Builder的构建参数,精度设置、工作空间、动态形状等。

    auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(
        createConfig(builder.get(), precision));
  6. 构建序列化网络:核心编译过程,执行各种优化:层融合、核函数选择、内存优化、布局转换等,生成可部署额序列化引擎文件,最耗时(约几分钟)最关键。

    auto serializedNetwork = std::unique_ptr<nvinfer1::IHostMemory>(
        builder->buildSerializedNetwork(*network, *config));
  7. 创建运行时(Runtime):Runtime是轻量级的引擎执行器,管理推理时的资源分配和执行。Runtime和Builder分离可以减少部署时的体积。

    runtime.reset(nvinfer1::createInferRuntime(*logger));
  8. 反序列化引擎:将序列化引擎文件转换为可执行引擎(序列化格式不能直接用于推理)。

    engine.reset(runtime->deserializeCudaEngine(serializedNetwork->data(),
        serializedNetwork->size()));
  9. 保存引擎文件:将构建的引擎持久化到磁盘,后续可直接加载进行推理。

    auto serialized = std::unique_ptr<nvinfer1::IHostMemory>(engine->serialize());
    	if (!serialized) {
    		std::cerr << "Failed to serialize engine" << std::endl;
    		return false;
    	}
    
    	std::ofstream file(enginePath, std::ios::binary);
    	if (!file.is_open()) {
    		std::cerr << "Failed to open file for writing: " << enginePath << std::endl;
    		return false;
    	}
    
    	file.write(reinterpret_cast<const char*>(serialized->data()), serialized->size());
    	file.close();

完整参考代码如下:

bool TrtEngine::buildFromOnnx(const std::string& onnxPath,
	const std::string& enginePath,
	nvinfer1::DataType precision) {
	std::cout << "Building TensorRT engine from ONNX: " << onnxPath << std::endl;

	// 创建构建器
	auto builder = std::unique_ptr<nvinfer1::IBuilder>(
		nvinfer1::createInferBuilder(*logger));
	if (!builder) {
		std::cerr << "Failed to create builder" << std::endl;
		return false;
	}

	// 创建网络定义
	const auto explicitBatch = 1U << static_cast<uint32_t>(
		nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
	auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
		builder->createNetworkV2(explicitBatch));
	if (!network) {
		std::cerr << "Failed to create network" << std::endl;
		return false;
	}

	// 创建解析器
	auto parser = std::unique_ptr<nvonnxparser::IParser>(
		nvonnxparser::createParser(*network, *logger));
	if (!parser) {
		std::cerr << "Failed to create parser" << std::endl;
		return false;
	}

	// 解析ONNX模型
	if (!parser->parseFromFile(onnxPath.c_str(),
		static_cast<int>(TrtLogger::Severity::kWARNING))) {
		std::cerr << "Failed to parse ONNX file" << std::endl;
		return false;
	}

	// 创建配置
	auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(
		createConfig(builder.get(), precision));
	if (!config) {
		std::cerr << "Failed to create config" << std::endl;
		return false;
	}

	// 构建序列化网络
	auto serializedNetwork = std::unique_ptr<nvinfer1::IHostMemory>(
		builder->buildSerializedNetwork(*network, *config));
	if (!serializedNetwork) {
		std::cerr << "Failed to build serialized network" << std::endl;
		return false;
	}

	// 创建运行时
	runtime.reset(nvinfer1::createInferRuntime(*logger));
	if (!runtime) {
		std::cerr << "Failed to create runtime" << std::endl;
		return false;
	}

	// 反序列化引擎
	engine.reset(runtime->deserializeCudaEngine(serializedNetwork->data(),
		serializedNetwork->size()));
	if (!engine) {
		std::cerr << "Failed to deserialize engine" << std::endl;
		return false;
	}

	// 保存引擎到文件
	if (!enginePath.empty()) {
		if (!saveEngine(enginePath)) {
			std::cerr << "Failed to save engine to file" << std::endl;
		}
	}

	std::cout << "Engine built successfully" << std::endl;
	return true;
}

bool TrtEngine::saveEngine(const std::string& enginePath) {
	if (!engine) {
		std::cerr << "No engine to save" << std::endl;
		return false;
	}

	auto serialized = std::unique_ptr<nvinfer1::IHostMemory>(engine->serialize());
	if (!serialized) {
		std::cerr << "Failed to serialize engine" << std::endl;
		return false;
	}

	std::ofstream file(enginePath, std::ios::binary);
	if (!file.is_open()) {
		std::cerr << "Failed to open file for writing: " << enginePath << std::endl;
		return false;
	}

	file.write(reinterpret_cast<const char*>(serialized->data()), serialized->size());
	file.close();

	std::cout << "Engine saved to: " << enginePath << std::endl;
	return true;
}

注: GUI工具Nsight Deep Learning Designer也可以构建引擎,相对与完成的TensorRT包来说更轻量级,如果仅有构建引擎的的需求的话,可以考虑此方式。Torch-TensorRT可以跳过ONNX导出环节直接生成引擎并完成推理,但适合用python进行快速验证,在工业C++部署场景仍然需要回到“PyTorch→ONNX→TensorRT C++” 流程。Torch-TensorRT使用示例如下:

import torch
import torch_tensorrt  # 导入Torch-TensorRT

# 1. 加载PyTorch模型(nnUNet模型)
model = torch.load("nnunet_model.pth").eval().cuda()
dummy_input = torch.randn(1,1,64,128,128).cuda()

# 2. 用Torch-TensorRT的高级API转换模型(核心:一行代码完成转换)
# compile方法就是“高级运行时API”,内部完成TensorRT引擎构建
trt_model = torch_tensorrt.compile(
    model,
    inputs=[dummy_input],  # 示例输入(用于推导维度)
    enabled_precisions={torch.float, torch.half},  # 支持FP32/FP16
    workspace_size=4 * 1024 * 1024 * 1024  # 工作空间大小
)

# 3. 推理:转换后的模型和普通PyTorch模型用法完全一致
output = trt_model(dummy_input)

三、加载引擎并推理

加载有引擎并推理有三种方式,分别是:

  1. 使用TensorRT 运行时 API。最常用的场景,本教程将以此为例。
  2. NVIDIA Dynamo-Triton。NVIDIA Dynamo-Triton(原名 NVIDIA Triton Inference Server)支持在包括 TensorRT、PyTorch、ONNX、OpenVINO、Python 和 RAPIDS FIL 在内的主要框架中部署 AI 模型。它通过动态批处理、并发执行和优化配置提供高性能。
  3. PyTorch 中部署。在PyTorch中部署即利用Torch-TensorRT直接生成引擎并推理,具体可见上一节介绍。

1、加载引擎

        加载引擎包括读取序列化的二进制引擎文件和反序列化引擎文件为可执行的引擎。参考代码如下:

bool TrtEngine::loadEngine(const std::string& enginePath) {
	std::cout << "Loading TensorRT engine from: " << enginePath << std::endl;

	// 读取引擎文件
	std::ifstream file(enginePath, std::ios::binary);
	if (!file.good()) {
		std::cerr << "Failed to open engine file" << std::endl;
		return false;
	}

	file.seekg(0, std::ios::end);
	size_t size = file.tellg();
	file.seekg(0, std::ios::beg);

	std::vector<char> engineData(size);
	file.read(engineData.data(), size);
	file.close();

	// 创建运行时
	runtime.reset(nvinfer1::createInferRuntime(*logger));
	if (!runtime) {
		std::cerr << "Failed to create runtime" << std::endl;
		return false;
	}

	// 反序列化引擎
	engine.reset(runtime->deserializeCudaEngine(engineData.data(), size));
	if (!engine) {
		std::cerr << "Failed to deserialize engine" << std::endl;
		return false;
	}

	std::cout << "Engine loaded successfully" << std::endl;
	return true;
}

2、执行推理

  1. 创建执行上下文。
    // 获取引擎
    nvinfer1::ICudaEngine* engine = buildEngine("model.onnx");
    
    // 创建执行上下文
    nvinfer1::IExecutionContext* context = engine->createExecutionContext();
  2. 准备输入输出缓冲区(分配设备内存)。

    // 获取输入输出张量信息
    int numIOTensors = engine->getNbIOTensors();
    std::vector<void*> deviceBuffers(numIOTensors);
    std::vector<size_t> bufferSizes(numIOTensors);
    
    // 为每个张量分配显存
    for (int i = 0; i < numIOTensors; ++i) {
        const char* name = engine->getIOTensorName(i);
        nvinfer1::Dims dims = engine->getTensorShape(name);
        nvinfer1::DataType dtype = engine->getTensorDataType(name);
        
        // 计算张量大小
        size_t vol = std::accumulate(dims.d, dims.d + dims.nbDims, 1, 
                                     std::multiplies<size_t>());
        size_t elementSize = 0;
        switch (dtype) {
            case nvinfer1::DataType::kFLOAT: elementSize = 4; break;
            case nvinfer1::DataType::kHALF: elementSize = 2; break;
            case nvinfer1::DataType::kINT8: elementSize = 1; break;
            default: elementSize = 4;
        }
        bufferSizes[i] = vol * elementSize;
        
        // 分配显存
        cudaMalloc(&deviceBuffers[i], bufferSizes[i]);
    }
  3. 设置输入数据。

    // 准备输入数据(以float为例)
    std::vector<float> inputData(1 * 3 * 224 * 224);  // batch=1, 3x224x224
    // 填充数据...
    
    // 找到输入张量索引
    int inputIndex = -1;
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        if (engine->getTensorIOMode(engine->getIOTensorName(i)) == 
            nvinfer1::TensorIOMode::kINPUT) {
            inputIndex = i;
            break;
        }
    }
    
    // 拷贝输入数据到显存
    cudaMemcpy(deviceBuffers[inputIndex], inputData.data(), 
               bufferSizes[inputIndex], cudaMemcpyHostToDevice);
  4. 执行推理。

    // 设置所有张量地址
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        const char* name = engine->getIOTensorName(i);
        context->setTensorAddress(name, deviceBuffers[i]);
    }
    
    // 创建CUDA流(可选)
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    
    // 同步推理
    context->enqueueV3(0);  // 使用默认流
    cudaStreamSynchronize(0);
    
    // 或异步推理
    context->enqueueV3(stream);
    // 可以在这里做其他CPU工作...
    cudaStreamSynchronize(stream);  // 最后同步
  5. 获取输出结果。

    // 找到输出张量
    int outputIndex = -1;
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        if (engine->getTensorIOMode(engine->getIOTensorName(i)) == 
            nvinfer1::TensorIOMode::kOUTPUT) {
            outputIndex = i;
            break;
        }
    }
    
    // 准备输出缓冲区
    std::vector<float> outputData(bufferSizes[outputIndex] / sizeof(float));
    
    // 从显存拷贝结果
    cudaMemcpy(outputData.data(), deviceBuffers[outputIndex], 
               bufferSizes[outputIndex], cudaMemcpyDeviceToHost);

    6、资源释放。

    // 释放显存
    for (void* buffer : deviceBuffers) {
        if (buffer) cudaFree(buffer);
    }
    
    // 释放CUDA流
    cudaStreamDestroy(stream);
    
    // 释放TensorRT资源
    delete context;
    delete engine;

        完成示例代码:
 

#include <iostream>
#include <vector>
#include <numeric>
#include <fstream>
#include <cuda_runtime.h>
#include <NvInfer.h>
#include <NvOnnxParser.h>

class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING) 
            std::cout << msg << std::endl;
    }
} gLogger;

int main() {
    // 1. 加载引擎
    auto runtime = std::unique_ptr<nvinfer1::IRuntime>(
        nvinfer1::createInferRuntime(gLogger));
    
    std::ifstream file("model.engine", std::ios::binary);
    file.seekg(0, std::ios::end);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    
    std::vector<char> buffer(size);
    file.read(buffer.data(), size);
    file.close();
    
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(
        buffer.data(), size);
    
    // 2. 创建上下文
    nvinfer1::IExecutionContext* context = engine->createExecutionContext();
    
    // 3. 准备缓冲区
    std::vector<void*> deviceBuffers(engine->getNbIOTensors());
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        const char* name = engine->getIOTensorName(i);
        auto dims = engine->getTensorShape(name);
        auto dtype = engine->getTensorDataType(name);
        
        size_t vol = std::accumulate(dims.d, dims.d + dims.nbDims, 1, 
                                     std::multiplies<size_t>());
        size_t elementSize = (dtype == nvinfer1::DataType::kFLOAT) ? 4 : 2;
        
        cudaMalloc(&deviceBuffers[i], vol * elementSize);
    }
    
    // 4. 设置输入数据(示例)
    std::vector<float> inputData(1*3*224*224, 1.0f);  // 全1输入
    int inputIdx = 0;  // 假设第一个张量是输入
    cudaMemcpy(deviceBuffers[inputIdx], inputData.data(), 
               inputData.size() * sizeof(float), cudaMemcpyHostToDevice);
    
    // 5. 设置张量地址
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        context->setTensorAddress(engine->getIOTensorName(i), deviceBuffers[i]);
    }
    
    // 6. 执行推理
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    context->enqueueV3(stream);
    cudaStreamSynchronize(stream);
    
    // 7. 获取输出
    int outputIdx = 1;  // 假设第二个张量是输出
    std::vector<float> outputData(1000);  // 假设输出大小为1000
    cudaMemcpy(outputData.data(), deviceBuffers[outputIdx], 
               outputData.size() * sizeof(float), cudaMemcpyDeviceToHost);
    
    // 8. 清理
    for (auto* buffer : deviceBuffers) cudaFree(buffer);
    cudaStreamDestroy(stream);
    delete context;
    delete engine;
    
    return 0;
}

参考:

https://docs.nvidia.com/deeplearning/tensorrt/10.14.1/getting-started/quick-start-guide.htmlhttps://docs.nvidia.com/deeplearning/tensorrt/10.14.1/getting-started/quick-start-guide.html

https://github.com/NVIDIA/TensorRT/blob/HEAD/quickstart/IntroNotebooks/2.%20Using%20PyTorch%20through%20ONNX.ipynbhttps://github.com/NVIDIA/TensorRT/blob/HEAD/quickstart/IntroNotebooks/2.%20Using%20PyTorch%20through%20ONNX.ipynb

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐