2. 语义分割模型基础
在进入TensorRT部署细节之前，我们必须先厘清语义分割模型在推理阶段产生的数据形态。很多部署工程师在转换ONNX或编写后处理代码时，常常因为混淆logits、概率图与类别索引而出现“输出错位”或“类别偏移”的问题。本章将系统梳理FCN、UNet、DeepLab三类经典网络的输出结构，明确从原始张量到最终掩码的完整数据链路，为后续TensorRT的显存分配和输出解析打下坚实基础。
2.1 核心原理：三类网络的输出形态与数据语义
语义分割网络本质上是一个逐像素分类器。无论网络结构如何变化，其最终输出都是一个包含类别信息的张量。理解输出格式的关键在于区分三个层次：logits（原始分数）、概率图（归一化分数）和类别索引（离散标签）。
以FCN（全卷积网络）为例，其输出层通常是一个1×1卷积，将特征图通道数映射为类别数C。假设输入图像尺寸为H×W，则输出张量形状为[1, C, H, W]。这里的每个通道对应一个类别，通道内每个像素位置的数值就是该类别在该位置的logit。logits可以是任意实数，正值表示该类别更有可能是该像素的标签，负值则相反。FCN的典型特点是使用转置卷积进行上采样，输出分辨率与输入相同，但边缘细节较粗糙。
UNet采用编码器-解码器结构，通过跳跃连接融合浅层与深层特征。其输出层同样产生[1, C, H, W]的logits张量。UNet的优势在于对小目标的分割更精细，但输出格式与FCN完全一致——都是未归一化的logits。DeepLab系列（如DeepLabV3+）引入了空洞卷积和ASPP（空洞空间金字塔池化）模块，输出层依然是一个卷积层，产生[1, C, H, W]的logits。值得注意的是，DeepLabV3+会先输出一个较低分辨率的logits（如输入尺寸的1/8），然后通过双线性插值上采样到原始分辨率。
从数学角度看，logits经过Softmax函数后得到概率图。Softmax的公式为：
p_i(x,y) = exp(logit_i(x,y)) / Σ_j exp(logit_j(x,y))其中i表示类别索引，x,y表示像素坐标。概率图的形状与logits相同，但每个通道的值被压缩到(0,1)区间，且所有通道在同一个像素位置的和为1。类别索引则是从logits或概率图中取最大值对应的通道编号，形状为[1, 1, H, W]或[1, H, W]，每个像素位置存储一个0到C-1的整数。
| 网络结构 | 输出层设计 | 输出张量形状 | 上采样方式 | 典型特点 |
|---|---|---|---|---|
| FCN | 1×1卷积 + 转置卷积 | [1, C, H, W] | 转置卷积（可学习） | 结构简单，感受野大，边缘粗糙 |
| UNet | 1×1卷积 | [1, C, H, W] | 转置卷积 + 跳跃连接 | 小目标分割效果好，特征融合充分 |
| DeepLabV3+ | 1×1卷积 + 双线性插值 | [1, C, H, W]（先低分辨率） | 双线性插值（固定） | ASPP多尺度特征，边界更精细 |
2.2 判断方法：如何从模型输出中正确提取所需数据
在实际部署中，我们需要根据下游任务决定使用哪一层输出。以下是三种常见场景的判断逻辑：
- 场景A：仅需最终分割掩码（如自动驾驶的障碍物区域提取）。此时直接对logits执行Argmax操作，得到类别索引图。不需要计算Softmax，因为Argmax在logits上的结果与在概率图上的结果完全一致（Softmax是单调递增函数）。这可以节省一次指数运算和除法运算。
- 场景B：需要置信度分数（如医学影像的病变概率评估）。此时必须对logits执行Softmax，得到每个类别的概率值。注意，Softmax的输出是C个通道，每个通道代表一个类别的概率，不能直接取最大值作为最终结果，而是需要结合业务逻辑（如设定阈值）。
- 场景C：需要多类别概率分布（如遥感地物分类中的不确定性分析）。此时需要完整的概率图，且可能需要保存为浮点类型（FP32或FP16），这会显著增加显存占用。
在TensorRT部署中，一个常见的优化技巧是在模型图中直接嵌入Argmax层，使网络输出直接为类别索引图，从而避免在GPU上单独执行Argmax操作。但这样做会丢失概率信息，因此需要根据业务需求权衡。另一个关键点是输出张量的内存布局：PyTorch默认使用NCHW格式，而TensorRT内部可能使用NC4HW2或NHWC等优化布局。在编写后处理代码时，必须明确当前张量的实际内存布局，否则按错误步长读取数据会导致类别错乱。
2.3 案例拆解：UNet医学影像分割的TensorRT输出解析
我们以一个实际的UNet肝脏分割项目为例，展示从模型输出到最终掩码的完整处理流程。该项目的输入是CT图像的512×512切片，类别数为2（背景和肝脏）。
背景与约束：医院要求推理延迟低于50ms，显存占用不超过2GB（与其它AI任务共享GPU）。初始方案直接在PyTorch中运行，延迟为120ms，无法满足要求。我们决定使用TensorRT进行加速，但需要处理输出格式的适配问题。
初始现象：将PyTorch导出的UNet模型转换为TensorRT引擎后，推理速度提升至35ms，但输出的掩码图像出现“棋盘格”状噪声，且肝脏区域边缘有1-2像素的偏移。
关键输入数据：PyTorch模型的输出为[1, 2, 512, 512]的logits张量，TensorRT引擎的输出binding形状相同。后处理代码使用torch.argmax(output, dim=1)获取类别索引。
分析判断过程：我们首先检查了TensorRT的输出数值，发现logits的数值范围与PyTorch一致，但内存布局不同。PyTorch的NCHW布局中，通道维是连续的；而TensorRT在FP16精度下可能采用NC4HW2布局，即每4个通道为一组，通道维被拆分。直接按NCHW的步长读取数据，会导致通道错位——原本的通道0数据被读成了通道1的数据，从而产生棋盘格噪声。
实际操作步骤：
- 在TensorRT的binding配置中，显式设置输出张量的格式为
TensorFormat.kLINEAR，强制使用NCHW线性布局。 - 在后处理代码中，使用
engine.get_binding_shape()动态获取输出形状，避免硬编码。 - 将Argmax操作从PyTorch迁移到CUDA kernel中，直接读取logits并输出类别索引，减少一次GPU→CPU的数据拷贝。
- 对于边缘偏移问题，我们对比了PyTorch和TensorRT的上采样层实现，发现TensorRT的双线性插值默认使用
align_corners=False，而PyTorch的UNet实现使用了align_corners=True。在ONNX导出时，该参数被固化，导致坐标映射差异。
结果对比：修复后，掩码图像与PyTorch输出完全一致，Dice系数从0.97提升至0.99（消除了棋盘格噪声），推理延迟稳定在32ms，显存占用1.6GB。
风险与代价：强制使用kLINEAR格式可能降低TensorRT的kernel优化空间（约5%的性能损失），但换来了正确的输出。如果业务允许，也可以在后处理中适配NC4HW2布局，但代码复杂度显著增加。
复盘经验：在部署任何分割模型前，务必先打印TensorRT引擎的输出张量形状和内存布局，与原始框架进行逐元素对比。不要假设输出格式与训练框架一致。
2.4 常见误区与总结复盘
在语义分割部署中，以下误区最为常见：
- 误区1：混淆logits与概率。有些工程师直接对logits执行阈值判断（如>0.5），这在二分类中看似合理，但多分类时完全错误。logits的值没有概率意义，必须经过Softmax或Sigmoid归一化。
- 误区2：忽略类别索引的数据类型。Argmax输出的类别索引通常是int64或int32，而可视化或存储时可能需要uint8。在GPU上直接转换数据类型会引入额外kernel，建议在TensorRT图中使用
Cast层完成转换。 - 误区3：在CPU上执行Argmax。如果输出张量是GPU显存，将其拷贝到CPU再执行Argmax会引入严重的PCIe传输延迟。正确做法是在GPU上完成所有后处理，只将最终的类别索引图（通常为uint8）拷贝回CPU。
- 误区4：忽略背景类别的处理。在计算mIoU或Dice时，背景类别（索引0）通常被排除，但部署时背景像素仍然占据大量显存。如果业务不需要背景，可以在后处理中直接过滤，但不要修改模型输出通道数。