TensorRT 10.3 不支持传统 Caffe 模型的解决方案与 GPU 设备优化实践

1次阅读

共计 2270 个字符，预计需要花费 6 分钟才能阅读完成。

TensorRT 作为 NVIDIA 推出的高性能深度学习推理引擎，在 10.3 版本中移除了对传统 Caffe 模型的支持。这一变化主要基于以下技术考量：

维护成本 ：Caffe 作为早期框架，社区活跃度下降，维护其解析器的成本过高
功能局限 ：Caffe 的 Layer 定义无法满足现代网络结构（如动态 shape、自定义算子）的需求
生态趋势 ：ONNX 已成为事实上的模型交换标准，支持 90%+ 的框架间转换

典型报错示例：

[TRT] Parameter check failed at: engine.cpp::nvinfer1::rt::CaffeParser::parse::1026

转换工具链 ：
caffe2onnx + onnx-simplifier 组合
支持 BatchNorm 等复杂结构的自动折叠
优势：
保留完整的模型结构和参数
支持后续的量化操作（INT8 校准）
ONNX Runtime 可作为验证中间件

PyTorch/TF 复现相同结构
使用原始数据 fine-tune
输出 ONNX/TensorRT 直接支持的格式

# 安装核心工具链（建议使用 docker 环境）!pip install onnx==1.14.0 onnxruntime-gpu==1.15.1 onnx-simplifier==0.4.33
!pip install torch==2.1.0 torchvision==0.16.0

import caffe2onnx

# 加载 caffe 模型
model = caffe2onnx.convert(
    prototxt_path="model.prototxt",
    caffemodel_path="model.caffemodel",
    output_path="model.onnx",
    opset_version=13  # 建议 >=11 以支持现代算子
)

# 简化 ONNX 模型
from onnxsim import simplify
simplified_model, check = simplify("model.onnx")
assert check, "Simplification failed"

import tensorrt as trt

# 创建 builder
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 解析 ONNX
parser = trt.OnnxParser(network, logger)
with open("model_simplified.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

# 配置优化参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 构建引擎
engine = builder.build_engine(network, config)
with open("model.engine", "wb") as f:
    f.write(engine.serialize())

显存池化 ：

config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 按需调整

动态形状优化 ：

profile = builder.create_optimization_profile()
profile.set_shape("input", (1,3,224,224), (8,3,224,224), (16,3,224,224))
config.add_optimization_profile(profile)

层融合 ：自动合并 Conv+BN+ReLU

精度控制 ：

config.set_flag(trt.BuilderFlag.FP16)  # 启用 FP16
# 或 INT8 校准
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator()  # 需实现校准数据接口