基于 TensorRT 的深度学习模型推理加速实战指南

// 构建阶段指定每个维度的范围
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("input", OptProfileSelector::kMIN, Dims4{1, 3, 256, 256});
profile->setDimensions("input", OptProfileSelector::kOPT, Dims4{1, 3, 640, 640});
profile->setDimensions("input", OptProfileSelector::kMAX, Dims4{1, 3, 1280, 1280});

# 拉取 TensorRT 官方镜像（选择和你的 CUDA 版本匹配的）
docker pull nvcr.io/nvidia/tensorrt:24.05-py3

# 启动容器
docker run --gpus all -it --rm \
  -v /your/workspace:/workspace \
  nvcr.io/nvidia/tensorrt:24.05-py3

# 解压
tar -xzf TensorRT-10.1.0.27.Ubuntu-22.04.x86_64-gnu.cuda-12.4.tar.gz

# 添加到环境变量
export TENSORRT_DIR=/path/to/TensorRT-10.1.0.27
export LD_LIBRARY_PATH=$TENSORRT_DIR/lib:$LD_LIBRARY_PATH
export PYTHONPATH=$TENSORRT_DIR/python:$PYTHONPATH

# 安装 Python 包
cd $TENSORRT_DIR/python
pip install tensorrt-10.1.0-cp310-none-linux_x86_64.whl

# 验证安装
python -c "import tensorrt; print(tensorrt.__version__)"

import tensorrt as trt
import torch

print(f"TensorRT version: {trt.__version__}")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device: {torch.cuda.get_device_name(0)}")

# 检查 TensorRT 核心库
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
print(f"TensorRT builder created successfully")

import torch
import torchvision.models as models

# 加载模型
model = models.resnet50(pretrained=True)
model.eval()
model.cuda()

# 构建 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224).cuda()

# 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    opset_version=17,           # 尽量用最新的 opset
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={              # 如果需要动态 shape
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    }
)

print("ONNX exported successfully")

# 安装 onnxsim
pip install onnxsim

# 简化模型
onnxsim resnet50.onnx resnet50_sim.onnx

# 或者用 Python API
from onnxsim import simplify
import onnx

model = onnx.load("resnet50.onnx")
model_sim, check = simplify(model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_sim, "resnet50_sim.onnx")

import tensorrt as trt

# 1. 创建 Logger
logger = trt.Logger(trt.Logger.WARNING)

# 2. 创建 Builder 和 Network
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 3. 创建 ONNX Parser
parser = trt.OnnxParser(network, logger)

# 4. 解析 ONNX 文件
success = parser.parse_from_file("resnet50_sim.onnx")
if not success:
    print("Failed to parse ONNX file")
    for error in parser.errors:
        print(error)
    exit(1)

# 5. 配置构建参数
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # 开启 FP16 精度
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB workspace

# 6. 构建序列化引擎
serialized_engine = builder.build_serialized_network(network, config)

# 7. 保存到文件
with open("resnet50_fp16.engine", "wb") as f:
    f.write(serialized_engine)

print("Engine built successfully!")

config.set_flag(trt.BuilderFlag.FP16)       # 开启 FP16
config.set_flag(trt.BuilderFlag.INT8)       # 开启 INT8
config.set_flag(trt.BuilderFlag.STRICT_TYPES) # 严格执行精度，不自动回退

config.set_flag(trt.BuilderFlag.DEBUG)      # 保留调试信息
config.set_flag(trt.BuilderFlag.PROFILING)  # 开启 profiling 层

config.set_flag(trt.BuilderFlag.TF32)       # 允许 TF32 计算（Ampere+）
config.set_flag(trt.BuilderFlag.FAST_MATH)  # 快速数学，可能有精度损失
config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS) # 优先保证精度

# 创建优化 profile
profile = builder.create_optimization_profile()

# 设置最小、最优、最大尺寸
profile.set_shape(
    "input",
    min=(1, 3, 224, 224),
    opt=(1, 3, 640, 640),
    max=(1, 3, 1280, 1280)
)

# 添加到 config
config.add_optimization_profile(profile)

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import os

class ImageBatchStream:
    def __init__(self, batch_size, calib_files):
        self.batch_size = batch_size
        self.calib_files = calib_files
        self.batch_count = len(calib_files) // batch_size
        self.max_batches = 100  # 用100个batch足够了
        
    def next_batch(self):
        for i in range(min(self.batch_count, self.max_batches)):
            batch = np.zeros((self.batch_size, 3, 224, 224), dtype=np.float32)
            for j in range(self.batch_size):
                img = self.load_image(self.calib_files[i * self.batch_size + j])
                batch[j] = self.preprocess(img)
            yield batch.ascontiguousarray()

class Int8Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, batch_stream, cache_file="calibration.cache"):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_stream = batch_stream
        self.cache_file = cache_file
        self.d_input = cuda.mem_alloc(4 * 3 * 224 * 224 * batch_stream.batch_size)
        self.batches = batch_stream.next_batch()
        
    def get_batch_size(self):
        return self.batch_stream.batch_size
    
    def get_batch(self, names):
        try:
            batch = next(self.batches)
            cuda.memcpy_htod(self.d_input, batch)
            return [int(self.d_input)]
        except StopIteration:
            return None
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None
    
    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

# 准备校准数据
calib_files = get_calibration_images("/path/to/coco/val2017", num_images=1000)
batch_stream = ImageBatchStream(batch_size=8, calib_files=calib_files)
calibrator = Int8Calibrator(batch_stream)

# 配置 INT8
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = calibrator

# 可以同时开启 FP16，TensorRT 会自动选择最优
config.set_flag(trt.BuilderFlag.FP16)

# 构建引擎
serialized_engine = builder.build_serialized_network(network, config)

inspector = engine.create_engine_inspector()
print(inspector.get_layer_information())

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

class TensorRTInfer:
    def __init__(self, engine_path):
        # 加载引擎
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 分配输入输出显存
        self.buffers = []
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding))
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            # 分配设备内存
            device_mem = cuda.mem_alloc(size * dtype().itemsize)
            self.buffers.append(device_mem)
        
        # 创建 CUDA 流
        self.stream = cuda.Stream()
        
    def infer(self, input_data):
        # input_data: numpy array on CPU
        # 1. 拷贝输入到 GPU
        cuda.memcpy_htod_async(self.buffers[0], input_data, self.stream)
        
        # 2. 执行推理
        self.context.execute_async_v2(
            bindings=[int(buf) for buf in self.buffers],
            stream_handle=self.stream.handle
        )
        
        # 3. 拷贝输出到 CPU
        output = np.empty(self.engine.get_binding_shape(1), dtype=np.float32)
        cuda.memcpy_dtoh_async(output, self.buffers[1], self.stream)
        
        # 4. 同步等待
        self.stream.synchronize()
        
        return output

# 初始化推理器
infer = TensorRTInfer("resnet50_fp16.engine")

# 预处理图片
image = load_image("test.jpg")
input_data = preprocess(image)  # shape (1, 3, 224, 224)

# 执行推理
output = infer.infer(input_data)

# 后处理
probabilities = softmax(output)
top5 = np.argsort(probabilities[0])[-5:][::-1]

print("Top-5 predictions:")
for idx in top5:
    print(f"  Class {idx}: {probabilities[0, idx]:.4f}")

import time
import torch
import torchvision.models as models

# PyTorch benchmark
model = models.resnet50(pretrained=True).cuda().half().eval()
dummy_input = torch.randn(1, 3, 224, 224).cuda().half()

# warmup
for _ in range(50):
    _ = model(dummy_input)
torch.cuda.synchronize()

# measure
start = time.time()
for _ in range(1000):
    _ = model(dummy_input)
torch.cuda.synchronize()
pytorch_time = (time.time() - start) / 1000 * 1000
print(f"PyTorch FP16: {pytorch_time:.2f} ms/image")

# TensorRT benchmark
infer = TensorRTInfer("resnet50_fp16.engine")
dummy_np = dummy_input.cpu().numpy()

# warmup
for _ in range(50):
    _ = infer.infer(dummy_np)

# measure
start = time.time()
for _ in range(1000):
    _ = infer.infer(dummy_np)
trt_time = (time.time() - start) / 1000 * 1000
print(f"TensorRT FP16: {trt_time:.2f} ms/image")
print(f"Speedup: {pytorch_time / trt_time:.2f}x")

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime_api.h>
#include <fstream>
#include <vector>
#include <memory>
#include <stdexcept>

class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING) {
            printf("[TensorRT] %s\n", msg);
        }
    }
};

class TensorRTInfer {
public:
    TensorRTInfer(const std::string& engine_path) {
        // 读取引擎文件
        std::ifstream file(engine_path, std::ios::binary);
        if (!file) {
            throw std::runtime_error("Cannot open engine file");
        }
        file.seekg(0, std::ios::end);
        size_t size = file.tellg();
        file.seekg(0, std::ios::beg);
        std::vector<char> engine_data(size);
        file.read(engine_data.data(), size);
        
        // 反序列化引擎
        m_runtime.reset(nvinfer1::createInferRuntime(m_logger));
        m_engine.reset(m_runtime->deserializeCudaEngine(
            engine_data.data(), size, nullptr
        ));
        if (!m_engine) {
            throw std::runtime_error("Failed to deserialize engine");
        }
        
        // 创建执行上下文
        m_context.reset(m_engine->createExecutionContext());
        
        // 分配显存
        m_buffers.resize(m_engine->getNbIOTensors());
        for (int i = 0; i < m_engine->getNbIOTensors(); i++) {
            const char* name = m_engine->getIOTensorName(i);
            auto dims = m_engine->getTensorShape(name);
            size_t bytes = nvinfer1::volume(dims) * sizeof(float);
            cudaMalloc(&m_buffers[i], bytes);
            
            if (m_engine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) {
                m_input_name = name;
                m_input_idx = i;
            } else {
                m_output_name = name;
                m_output_idx = i;
            }
        }
        
        // 创建 CUDA 流
        cudaStreamCreate(&m_stream);
    }
    
    ~TensorRTInfer() {
        for (auto buf : m_buffers) {
            cudaFree(buf);
        }
        cudaStreamDestroy(m_stream);
    }
    
    // 禁用拷贝
    TensorRTInfer(const TensorRTInfer&) = delete;
    TensorRTInfer& operator=(const TensorRTInfer&) = delete;
    
    void infer(const float* input, float* output, int batch_size = 1) {
        // 设置输入 shape（如果是动态的）
        auto dims = m_engine->getTensorShape(m_input_name);
        if (dims.d[0] == -1) {
            dims.d[0] = batch_size;
            m_context->setInputShape(m_input_name, dims);
        }
        
        // H2D 拷贝
        cudaMemcpyAsync(
            m_buffers[m_input_idx], input,
            nvinfer1::volume(dims) * sizeof(float),
            cudaMemcpyHostToDevice, m_stream
        );
        
        // 设置张量地址
        m_context->setTensorAddress(m_input_name, m_buffers[m_input_idx]);
        m_context->setTensorAddress(m_output_name, m_buffers[m_output_idx]);
        
        // 执行推理
        m_context->enqueueV3(m_stream);
        
        // D2H 拷贝
        auto out_dims = m_context->getTensorShape(m_output_name);
        cudaMemcpyAsync(
            output, m_buffers[m_output_idx],
            nvinfer1::volume(out_dims) * sizeof(float),
            cudaMemcpyDeviceToHost, m_stream
        );
        
        // 同步
        cudaStreamSynchronize(m_stream);
    }
    
private:
    Logger m_logger;
    std::unique_ptr<nvinfer1::IRuntime> m_runtime;
    std::unique_ptr<nvinfer1::ICudaEngine> m_engine;
    std::unique_ptr<nvinfer1::IExecutionContext> m_context;
    std::vector<void*> m_buffers;
    cudaStream_t m_stream;
    const char* m_input_name;
    const char* m_output_name;
    int m_input_idx;
    int m_output_idx;
};

cmake_minimum_required(VERSION 3.18)
project(tensorrt_infer)

set(CMAKE_CXX_STANDARD 17)

# CUDA
find_package(CUDA REQUIRED)
include_directories(${CUDA_INCLUDE_DIRS})

# TensorRT
set(TENSORRT_ROOT /path/to/TensorRT-10.1.0.27)
include_directories(${TENSORRT_ROOT}/include)
link_directories(${TENSORRT_ROOT}/lib)

# 可执行文件
add_executable(infer main.cpp)
target_link_libraries(infer
    ${CUDA_LIBRARIES}
    nvinfer
    nvonnxparser
    cudart
)

int main() {
    try {
        TensorRTInfer infer("resnet50_fp16.engine");
        
        // 准备输入
        std::vector<float> input(1 * 3 * 224 * 224);
        std::vector<float> output(1 * 1000);
        
        // 填充 input...
        
        // 执行推理
        infer.infer(input.data(), output.data());
        
        // 处理 output...
        
        std::cout << "Inference done!" << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    return 0;
}

# 创建多个推理实例，每个实例有自己的 stream
infer1 = TensorRTInfer("model.engine")
infer2 = TensorRTInfer("model.engine")

# 在不同的线程中跑各自的推理
# 它们会在 GPU 上并发执行

config.set_flag(trt.BuilderFlag.STRIP_PLAN)