4. TensorRT深度解析（下）

上一讲我们把TensorRT的静态图优化聊透了。今天咱们来啃几块硬骨头——动态形状、自定义插件，还有API实战。说实话，这三块才是真正区分“会用”和“用好”的分水岭。

我在生产环境里踩过的坑，多半都跟动态形状和插件有关。你想想看，模型部署到线上，输入尺寸不可能永远固定。客户传一张1920×1080的图，你总不能让人家先缩放到512×512吧？

4.1 动态形状支持：别再写死输入尺寸了

TensorRT从7.0开始支持动态形状。说白了，就是允许你在推理时改变输入的batch size、通道数、高和宽。但代价是——性能会打折扣，因为优化器没法做极致的常量折叠。

4.1.1 动态形状的三要素

要启用动态形状，你得在构建引擎时指定三个范围：

我个人习惯把opt设成线上最常见的batch size。比如线上QPS要求高，batch=8最常用，那就把opt设成8。TensorRT会针对这个尺寸做最激进的优化。

4.1.2 C++ API配置动态形状

// 创建builder和network
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U);

// 定义输入张量，维度用-1表示动态
nvinfer1::ITensor* input = network->addInput(
    "input", nvinfer1::DataType::kFLOAT, 
    nvinfer1::Dims4{-1, 3, -1, -1});

// 创建优化配置
nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("input", nvinfer1::OptProfileSelector::kMIN, 
                       nvinfer1::Dims4{1, 3, 224, 224});
profile->setDimensions("input", nvinfer1::OptProfileSelector::kOPT, 
                       nvinfer1::Dims4{8, 3, 224, 224});
profile->setDimensions("input", nvinfer1::OptProfileSelector::kMAX, 
                       nvinfer1::Dims4{32, 3, 224, 224});

// 配置config
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
config->addOptimizationProfile(profile);

// 构建引擎
nvinfer1::IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);

这段代码里，Dims4{-1, 3, -1, -1}表示batch和宽高都是动态的。注意通道数我写死了3——除非你处理的是多模态数据，否则通道数一般固定。

4.1.3 推理时设置输入尺寸

引擎构建好之后，每次推理前你得告诉它这次用多大尺寸：

// 创建执行上下文
nvinfer1::IExecutionContext* context = engine->createExecutionContext();

// 假设这次输入是 batch=4, 640x480
nvinfer1::Dims4 inputDims{4, 3, 480, 640};
context->setBindingDimensions(0, inputDims);

// 检查是否有效
if (!context->allInputDimensionsSpecified()) {
    // 处理错误
}

// 执行推理
context->enqueueV2(buffers, stream, nullptr);

4.2 插件编写：当TensorRT不够用时

TensorRT内置了200多种算子，但总有些骚操作它不支持。比如我做过一个模型，里面有个自定义的“可变形卷积”，TensorRT压根不认识。这时候就得自己写插件。

插件说白了，就是告诉TensorRT：“这个算子的计算逻辑是什么，梯度怎么传，内存怎么分配”。

4.2.1 插件基类结构

TensorRT的插件体系分两层：

• IPluginV2DynamicExt：支持动态形状的插件基类（推荐）
• IPluginV2IOExt：旧版基类，不支持动态形状

我建议你直接用IPluginV2DynamicExt。虽然接口多一点，但兼容性更好。下面是一个简单的ReLU插件骨架：

class MyReLUPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
    // 返回插件名称
    const char* getPluginType() const noexcept override {
        return "MyReLU";
    }
    
    // 返回插件版本
    const char* getPluginVersion() const noexcept override {
        return "1";
    }
    
    // 获取输出维度（动态形状下需要根据输入计算）
    nvinfer1::DimsExprs getOutputDimensions(
        int outputIndex, const nvinfer1::DimsExprs* inputs, 
        int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override {
        // ReLU不改变形状，直接返回输入维度
        return inputs[0];
    }
    
    // 检查输入输出格式是否支持
    bool supportsFormatCombination(
        int pos, const nvinfer1::PluginTensorDesc* inOut, 
        int nbInputs, int nbOutputs) noexcept override {
        // 只支持float32
        return inOut[pos].type == nvinfer1::DataType::kFLOAT;
    }
    
    // 配置插件（分配workspace等）
    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 kernel调用
        const float* input = static_cast<const float*>(inputs[0]);
        float* output = static_cast<float*>(outputs[0]);
        int64_t volume = 1;
        for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) {
            volume *= inputDesc[0].dims.d[i];
        }
        myReluKernel<<<(volume + 255) / 256, 256, 0, stream>>>(input, output, volume);
        return 0;
    }
    
    // 克隆插件
    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override {
        return new MyReLUPlugin(*this);
    }
    
    // 序列化
    void serialize(void* buffer) const noexcept override {
        // 如果有成员变量，需要序列化
    }
    
    // 反序列化构造
    MyReLUPlugin(const void* data, size_t length) {
        // 从buffer恢复成员变量
    }
};

4.2.2 插件注册与使用

写好了插件类，还得注册到TensorRT的插件工厂里：

// 注册插件
class MyPluginFactory : public nvinfer1::IPluginFactory {
public:
    nvinfer1::IPluginV2* createPlugin(
        const char* layerName, const void* serialData, 
        size_t serialLength) noexcept override {
        return new MyReLUPlugin(serialData, serialLength);
    }
};

// 反序列化引擎时传入工厂
MyPluginFactory factory;
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(
    modelData, modelSize, &factory);

在ONNX模型中，你可以通过自定义节点来触发插件。比如在PyTorch中导出ONNX时，用torch.onnx.register_custom_op把自定义操作映射到插件名称上。

4.3 C++/Python API实战：两种语言，一个引擎

TensorRT提供了C++和Python两套API。C++性能更好，Python开发更快。我的建议是：原型验证用Python，生产部署用C++。

4.3.1 Python API快速上手

Python API最大的优势是跟PyTorch/TensorFlow无缝衔接。下面是一个完整的推理流程：

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

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

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

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

# 分配GPU内存
inputs = []
outputs = []
bindings = []
for i in range(engine.num_bindings):
    shape = engine.get_binding_shape(i)
    size = trt.volume(shape)
    dtype = trt.nptype(engine.get_binding_dtype(i))
    
    # 分配host和device内存
    host_mem = cuda.pagelocked_empty(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)
    
    bindings.append(int(device_mem))
    if engine.binding_is_input(i):
        inputs.append({"host": host_mem, "device": device_mem})
    else:
        outputs.append({"host": host_mem, "device": device_mem})

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
np.copyto(inputs[0]["host"], input_data.ravel())

# 推理
stream = cuda.Stream()
cuda.memcpy_htod_async(inputs[0]["device"], inputs[0]["host"], stream)
context.execute_async_v2(bindings, stream.handle)
cuda.memcpy_dtoh_async(outputs[0]["host"], outputs[0]["device"], stream)
stream.synchronize()

# 获取结果
result = outputs[0]["host"].reshape(engine.get_binding_shape(1))
print(result)

这段代码里，cuda.pagelocked_empty分配的是页锁定内存，能加速CPU到GPU的传输。我刚开始用普通numpy数组，结果传输速度慢了一倍。

4.3.2 C++ API性能调优

C++ API在性能调优上更灵活。比如你可以用enqueueV2配合CUDA stream实现异步推理：

// 创建多个stream实现并发
cudaStream_t streams[4];
for (int i = 0; i < 4; ++i) {
    cudaStreamCreate(&streams[i]);
}

// 每个stream绑定一个context
nvinfer1::IExecutionContext* contexts[4];
for (int i = 0; i < 4; ++i) {
    contexts[i] = engine->createExecutionContext();
}

// 并发推理
for (int i = 0; i < 4; ++i) {
    contexts[i]->enqueueV2(buffers[i], streams[i], nullptr);
}

// 同步所有stream
for (int i = 0; i < 4; ++i) {
    cudaStreamSynchronize(streams[i]);
}

这种多stream并发的方式，能把GPU利用率从60%拉到95%以上。我在一个OCR服务里就是这么干的——4个stream处理4张图片，吞吐量直接翻倍。

4.3.3 性能对比：Python vs C++

我个人的经验是：延迟敏感场景用C++，吞吐量敏感场景用Python。为什么？因为Python的GIL在多batch推理时影响不大，但单batch延迟会多出1-2ms。

4.4 本章知识体系

下面这张图总结了动态形状、插件编写和API实战的核心脉络：