网络定义与解析：使用TensorRT API构建网络，支持ONNX、Caffe等格式的解析

TensorRT的工作流可以概括为三个阶段：网络定义、引擎构建与运行时推理。其中网络定义是后续所有优化的基础——它决定了TensorRT能看到哪些层、哪些张量、哪些数据流关系。如果网络定义阶段出现偏差，后续的算子融合、精度校准、内核自动调优都将失去意义。本章聚焦于如何通过TensorRT API构建网络，并解析ONNX、Caffe等主流模型格式。

4.1 核心原理：从框架模型到TensorRT网络

TensorRT并不直接读取PyTorch或TensorFlow的权重文件，而是通过中间表示（IR）来接收模型。当前最主流的中间表示是ONNX（Open Neural Network Exchange），它充当了框架与TensorRT之间的“通用语言”。Caffe模型则通过TensorRT自带的caffeParser解析。无论是哪种格式，最终都会被转换为TensorRT内部的INetworkDefinition对象。

理解这一过程的关键在于：解析器（Parser）负责将外部格式的算子映射为TensorRT的层（Layer），张量（Tensor）则映射为ITensor。映射过程中，解析器会进行初步的算子识别与维度推断，但不会执行任何融合或精度调整——这些是后续构建引擎时由优化器完成的。

核心概念速览：

INetworkDefinition：TensorRT网络定义的顶层容器，包含层、张量、输入输出标记。
ONNXParser：解析ONNX模型，支持opset版本范围较广，是当前最常用的入口。
CaffeParser：解析Caffe的.prototxt与.caffemodel，适用于旧版模型。
UFF（已废弃）：曾用于TensorFlow，现已不再推荐，建议通过ONNX转换。

下图展示了从框架模型到TensorRT网络定义的核心数据流，以及解析器在其中扮演的桥梁角色：

4.2 方法流程：ONNX解析的完整步骤

使用ONNX解析器构建网络，是当前最主流的路径。其核心流程可以归纳为以下五个步骤：

创建Builder与Network：通过createInferBuilder获取builder，再通过builder.createNetworkV2(0)创建网络定义。其中参数0表示使用显式批处理维度（explicit batch），这是现代TensorRT（8.x及以上）的推荐方式。
创建ONNX解析器：调用nvinfer1::createParser，并将网络定义与builder的日志记录器传入。
解析ONNX文件：调用parser->parseFromFile(model_path, verbosity)。解析成功后，网络定义中会填充所有层与张量。
标记输入与输出：通过network->getInput(0)和network->getOutput(0)获取输入输出张量，并设置名称、数据类型、维度等属性。
验证网络：检查层数、张量维度是否符合预期，必要时使用network->validate()进行一致性校验。

以下代码展示了ONNX解析的核心片段：

// 创建 builder 与 network
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));

// 创建 ONNX 解析器
auto parser = std::unique_ptr<nvinfer1::IONNXParser>(nvinfer1::createParser(*network, logger));

// 解析 ONNX 文件
if (!parser->parseFromFile("resnet50.onnx", static_cast<int>(nvinfer1::ILogger::Severity::kWARNING))) {
    std::cerr << "ERROR: 解析 ONNX 文件失败" << std::endl;
    return -1;
}

// 标记输入输出
auto input = network->getInput(0);
input->setName("input_tensor");
input->setType(nvinfer1::DataType::kFLOAT);
input->setDimensions(nvinfer1::Dims4{1, 3, 224, 224});

auto output = network->getOutput(0);
output->setName("output_tensor");

关键提示：在TensorRT 8.x中，createNetworkV2必须传入kEXPLICIT_BATCH标志，否则无法处理动态维度。旧版createNetwork已废弃，请勿在新项目中使用。

4.3 案例：Caffe模型解析与自定义层处理

尽管ONNX已成为主流，但仍有大量遗留的Caffe模型需要部署。TensorRT提供了caffeParser来解析Caffe的.prototxt与.caffemodel文件。其使用方式与ONNX解析器类似，但需要注意以下几点：

Caffe的BatchNorm层在TensorRT中通常会被融合进卷积层，解析器会自动处理。
Caffe的Scale层可能被识别为逐元素乘法，需要确认其与BatchNorm的配合。
如果Caffe模型中包含TensorRT不支持的层（如自定义Python层），解析会失败，此时需要手动实现IPluginV2扩展。

以下是一个Caffe解析的典型代码结构：

auto caffeParser = std::unique_ptr<nvinfer1::ICaffeParser>(nvinfer1::createCaffeParser());

// 解析 prototxt 与 caffemodel
const auto blobNameToTensor = caffeParser->parse("deploy.prototxt", "model.caffemodel", *network, nvinfer1::DataType::kFLOAT);

// 获取输入输出张量
auto input = blobNameToTensor->find("data");
auto output = blobNameToTensor->find("prob");

// 标记输入输出
network->markOutput(*output);

在实际项目中，我们曾遇到一个Caffe的Eltwise层（元素相加）在TensorRT中被错误识别为拼接层的情况。排查后发现，Caffe的Eltwise参数中operation字段为SUM，但TensorRT的caffeParser在旧版本中对该字段的解析存在歧义。解决方案是手动修改.prototxt，将Eltwise替换为两个独立的张量相加操作，或者升级到TensorRT 8.5+版本，该问题已修复。

4.4 误区与复盘：网络定义阶段的常见陷阱

网络定义阶段看似简单，但实践中经常出现以下问题，导致后续引擎构建失败或推理结果错误：

误区	现象	根因	解决方案
忽略显式批处理标志	动态维度无法设置	使用了旧版`createNetwork`	使用`createNetworkV2(kEXPLICIT_BATCH)`
输入维度与ONNX不一致	解析成功但推理崩溃	ONNX中维度为动态，未指定具体值	在标记输入时显式设置`Dims4`
未处理不支持的自定义层	解析器返回错误	模型包含TensorRT不认识的算子	实现`IPluginV2`或修改模型结构
输出张量未标记	引擎构建成功但无法获取结果	忘记调用`markOutput`	检查所有输出张量均已标记
权重精度不匹配	精度下降明显	模型为FP32，但网络设置为FP16	在builder配置中设置`setFlag(kFP16)`前确认

重要提醒：网络定义阶段不会执行任何算子融合或精度校准。如果你在解析后立即检查层数，会发现与原始模型几乎一致。真正的优化发生在builder->buildEngineWithConfig()阶段。因此，不要试图在网络定义阶段手动“优化”网络结构——这反而可能破坏TensorRT的自动优化能力。

复盘一个真实案例：某团队在部署YOLOv5模型时，ONNX解析成功，但推理结果全为0。排查后发现，他们在标记输入时使用了Dims3{3, 640, 640}，而ONNX模型中实际是Dims4{1, 3, 640, 640}。由于显式批处理模式下维度必须完全匹配，导致输入数据被错误解释。修正为Dims4后问题解决。这个案例说明，维度一致性是网络定义阶段最容易忽略但影响最大的细节。

另一个常见问题是动态形状（Dynamic Shape）的处理。如果ONNX模型包含动态维度（如batch=-1），必须在网络定义阶段通过network->setInputShape设置优化轮廓（optimization profile），否则引擎构建会失败。建议在解析前先使用onnx.shape_inference工具检查模型的静态维度情况。

最后，建议在解析完成后打印网络结构摘要，确认层数与关键张量维度符合预期。可以使用以下代码快速验证：

std::cout << "网络层数: " << network->getNbLayers() << std::endl;
for (int i = 0; i < network->getNbLayers(); ++i) {
    auto layer = network->getLayer(i);
    std::cout << "层 " << i << ": " << layer->getName()
              << " (" << static_cast<int>(layer->getType()) << ")" << std::endl;
}

网络定义是整个TensorRT流程的基石。只有确保解析正确、维度一致、输入输出标记完整，后续的引擎构建与推理才能稳定高效。在实际项目中，建议将网络定义封装为独立的模块，并加入日志输出与断言检查，以便在模型更新时快速定位问题。