⚡

TensorRT模型优化与部署Skill tensorrt-optimization

TensorRT模型优化与部署技能是一个专注于利用NVIDIA TensorRT技术，将深度学习模型（如ONNX格式）高效转换为优化后的推理引擎的专业工具。它支持FP16、INT8等多种精度模式，提供动态形状处理、内核融合分析、自定义插件开发以及详细的性能分析功能。通过该技能，开发者可以显著提升模型在GPU上的推理速度与效率，降低延迟，实现高性能AI应用部署。关键词：TensorRT，模型优化，推理加速，INT8校准，动态形状，CUDA，深度学习部署，性能分析。

深度学习 1 次安装 12 次浏览 更新于 2/25/2026

name: tensorrt-optimization description: NVIDIA TensorRT 模型优化与部署。将模型转换为 TensorRT 引擎，配置优化配置文件和精度模式，应用 INT8 校准，分析内核融合，生成自定义插件，并分析推理性能。 allowed-tools: Bash(*) Read Write Edit Glob Grep WebFetch metadata: author: babysitter-sdk version: “1.0.0” category: ml-inference backlog-id: SK-008

tensorrt-optimization

您是 tensorrt-optimization - 一个专门用于 NVIDIA TensorRT 模型优化和部署的技能。此技能提供了优化深度学习模型以进行推理的专家级能力。

概述

此技能支持 AI 驱动的 TensorRT 优化，包括：

  • 将模型转换为 TensorRT 引擎
  • 配置优化配置文件和精度模式
  • 应用 INT8 校准和量化
  • 分析内核融合机会
  • 生成自定义 TensorRT 插件
  • 分析推理延迟和吞吐量
  • 处理动态形状和批处理大小
  • 比较 TensorRT 与框架推理

先决条件

  • TensorRT 8.5+
  • CUDA Toolkit 11.0+
  • ONNX Runtime（用于 ONNX 模型）
  • Python TensorRT 包

能力

1. 模型转换为 TensorRT

从各种框架转换模型：

import tensorrt as trt

# 创建构建器和网络
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 解析 ONNX 模型
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 配置构建器
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

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

# 保存引擎
with open("model.engine", "wb") as f:
    f.write(engine)

2. 精度配置

配置 FP16、INT8 和 TF32：

# 启用 FP16
config.set_flag(trt.BuilderFlag.FP16)

# 启用 INT8（需要校准）
config.set_flag(trt.BuilderFlag.INT8)

# 启用 TF32（Ampere+）
config.clear_flag(trt.BuilderFlag.TF32)  # 如果需要可以禁用

# 启用稀疏张量核心
config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)

# 按层偏好精度
config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)

# 强制严格类型
config.set_flag(trt.BuilderFlag.STRICT_TYPES)

3. INT8 校准

class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file):
        super().__init__()
        self.data_loader = iter(data_loader)
        self.cache_file = cache_file
        self.batch_size = data_loader.batch_size
        self.device_input = cuda.mem_alloc(
            self.batch_size * 3 * 224 * 224 * 4)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        try:
            batch = next(self.data_loader)
            cuda.memcpy_htod(self.device_input, batch.numpy())
            return [int(self.device_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)

# 使用校准器
calibrator = Calibrator(calibration_loader, "calibration.cache")
config.int8_calibrator = calibrator
config.set_flag(trt.BuilderFlag.INT8)

4. 动态形状

处理可变输入大小：

# 创建优化配置文件
profile = builder.create_optimization_profile()

# 定义形状范围 [最小, 最优, 最大]
profile.set_shape("input",
    min=(1, 3, 224, 224),     # 最小形状
    opt=(8, 3, 224, 224),     # 最优形状
    max=(32, 3, 224, 224))    # 最大形状

config.add_optimization_profile(profile)

# 针对不同场景的多个配置文件
profile_small = builder.create_optimization_profile()
profile_small.set_shape("input", (1, 3, 224, 224), (4, 3, 224, 224), (8, 3, 224, 224))
config.add_optimization_profile(profile_small)

profile_large = builder.create_optimization_profile()
profile_large.set_shape("input", (16, 3, 224, 224), (32, 3, 224, 224), (64, 3, 224, 224))
config.add_optimization_profile(profile_large)

5. 推理执行

# 加载引擎
runtime = trt.Runtime(logger)
with open("model.engine", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

# 创建执行上下文
context = engine.create_execution_context()

# 为动态形状设置输入形状
context.set_input_shape("input", (batch_size, 3, 224, 224))

# 分配缓冲区
inputs = []
outputs = []
bindings = []

for i in range(engine.num_io_tensors):
    name = engine.get_tensor_name(i)
    dtype = trt.nptype(engine.get_tensor_dtype(name))
    shape = context.get_tensor_shape(name)
    size = trt.volume(shape)

    buffer = cuda.mem_alloc(size * dtype.itemsize)
    bindings.append(int(buffer))

    if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
        inputs.append(buffer)
    else:
        outputs.append(buffer)

# 执行推理
cuda.memcpy_htod(inputs[0], input_data)
context.execute_v2(bindings)
cuda.memcpy_dtoh(output_data, outputs[0])

6. 插件开发

创建自定义操作：

// 插件类
class CustomPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
    int getNbOutputs() const noexcept override { return 1; }

    nvinfer1::DimsExprs getOutputDimensions(
        int outputIndex,
        const nvinfer1::DimsExprs* inputs,
        int nbInputs,
        nvinfer1::IExprBuilder& exprBuilder) noexcept override {
        return inputs[0];  // 与输入形状相同
    }

    int enqueue(
        const nvinfer1::PluginTensorDesc* inputDesc,
        const nvinfer1::PluginTensorDesc* outputDesc,
        const void* const* inputs,
        void* const* outputs,
        void* workspace,
        cudaStream_t stream) noexcept override {
        // 启动自定义 CUDA 内核
        customKernel<<<blocks, threads, 0, stream>>>(
            inputs[0], outputs[0], inputDesc[0].dims);
        return 0;
    }
};

// 注册插件
REGISTER_TENSORRT_PLUGIN(CustomPluginCreator);

7. 性能分析

# 启用分析
config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED

# 使用时序缓存以加速构建
timing_cache_file = "timing.cache"
if os.path.exists(timing_cache_file):
    with open(timing_cache_file, "rb") as f:
        cache = config.create_timing_cache(f.read())
else:
    cache = config.create_timing_cache(b"")
config.set_timing_cache(cache, ignore_mismatch=False)

# 分析推理
profiler = trt.Profiler()
context.profiler = profiler

# 基准测试
import time
warmup = 10
iterations = 100

for _ in range(warmup):
    context.execute_v2(bindings)
cuda.Context.synchronize()

start = time.perf_counter()
for _ in range(iterations):
    context.execute_v2(bindings)
cuda.Context.synchronize()
end = time.perf_counter()

latency = (end - start) / iterations * 1000
throughput = batch_size * iterations / (end - start)
print(f"延迟: {latency:.2f} ms, 吞吐量: {throughput:.2f} 样本/秒")

8. 内核融合分析

# 使用 trtexec 进行分析
trtexec --onnx=model.onnx \
    --fp16 \
    --workspace=4096 \
    --verbose \
    --dumpLayerInfo \
    --exportLayerInfo=layers.json

# 使用 Nsight Systems 进行分析
nsys profile -o trt_profile \
    trtexec --loadEngine=model.engine --iterations=100

# 查看层时序
trtexec --loadEngine=model.engine \
    --dumpProfile \
    --separateProfileRun

命令行工具

# 将 ONNX 转换为 TensorRT
trtexec --onnx=model.onnx --saveEngine=model.engine

# 使用 FP16
trtexec --onnx=model.onnx --fp16 --saveEngine=model_fp16.engine

# 使用 INT8 校准
trtexec --onnx=model.onnx --int8 \
    --calib=calibration.cache --saveEngine=model_int8.engine

# 动态形状
trtexec --onnx=model.onnx \
    --minShapes=input:1x3x224x224 \
    --optShapes=input:8x3x224x224 \
    --maxShapes=input:32x3x224x224 \
    --saveEngine=model_dynamic.engine

# 基准测试现有引擎
trtexec --loadEngine=model.engine \
    --iterations=1000 \
    --warmUp=500 \
    --duration=10

流程集成

此技能与以下流程集成：

  • ml-inference-optimization.js - ML 推理优化
  • tensor-core-programming.js - 张量核心使用

输出格式

{
  "operation": "build-engine",
  "status": "success",
  "input_model": "model.onnx",
  "output_engine": "model.engine",
  "configuration": {
    "precision": ["FP16", "INT8"],
    "workspace_mb": 1024,
    "dynamic_shapes": true
  },
  "optimization": {
    "layer_fusions": 23,
    "reformats_eliminated": 8,
    "tactics_selected": 156
  },
  "performance": {
    "build_time_s": 45.2,
    "engine_size_mb": 28.5,
    "estimated_latency_ms": 1.2
  }
}

依赖项

  • TensorRT 8.5+
  • CUDA Toolkit 11.0+
  • ONNX Runtime（可选）
  • Python tensorrt 包

约束

  • INT8 需要代表性校准数据
  • 动态形状会增加构建时间
  • 自定义插件需要仔细的内存管理
  • 引擎文件是 GPU 架构特定的