ONNX 转 TensorRT Bug 记录：IIfConditionalOutputLayer

1. 问题描述

环境：TensorRT-8.6.1.6、CUDA-11.8

报错：Error[4]: /If_OutputLayer: IIfConditionalOutputLayer inputs must have the same shape. Shapes are [-1,384] and [-1,1,384].

复现代码：

python">import os
import torch
import torch.nn as nn
import torch.nn.functional as Fclass TestModel(torch.nn.Module):def __init__(self, mode):super().__init__()self.mode = modeself.conv = nn.Conv1d(512, 512, 3, 2, 1)def forward(self, x, mask):if self.mode == 1:return self.forward1(x, mask)elif self.mode == 2:return self.forward2(x, mask)elif self.mode == 3:return self.forward3(x, mask)elif self.mode == 4:return self.forward4(x, mask)else:raise ValueError("Invalid mode")def forward1(self, x, mask):mask = mask.unsqueeze(1)x = self.conv(x)mask = F.interpolate(mask, size=x.size(-1), mode="nearest")x = x * maskmask = mask.squeeze(1)return x, maskdef forward2(self, x, mask):mask = mask.unsqueeze(1)x = self.conv(x)mask = F.interpolate(mask, size=384, mode="nearest")x = x * maskmask = mask.squeeze(1)return x, maskdef forward3(self, x, mask):mask = mask.unsqueeze(1)x = x * maskmask = mask.squeeze(1)return x, maskdef forward4(self, x, mask):mask = mask.unsqueeze(1)x = self.conv(x)mask = F.interpolate(mask, size=x.size(-1), mode="nearest")x = x * maskb = x.shape[0]mask = mask.reshape(b, -1)return x, maskfake_input = torch.randn(1, 512, 768)
fake_mask = torch.randn(1, 768)
model1 = TestModel(1)
model2 = TestModel(2)
model3 = TestModel(3)
model4 = TestModel(4)with torch.no_grad():print([x.shape for x in model1(fake_input, fake_mask)])print([x.shape for x in model2(fake_input, fake_mask)])print([x.shape for x in model3(fake_input, fake_mask)])print([x.shape for x in model4(fake_input, fake_mask)])dynamic = {"x_input": {0: "batch"},"masks_input": {0: "batch"},"x_output": {0: "batch"},"masks_output": {0: "batch"}
}save_dir = "log"
os.makedirs(save_dir, exist_ok=True)for i, model in enumerate((model1, model2, model3, model4)):torch.onnx.export(model.cpu().eval(),(fake_input.cpu(), fake_mask.cpu()),os.path.join(save_dir, f'dynamic_{i+1}.onnx'),verbose=True,opset_version=17,do_constant_folding=True,input_names=["x_input", "masks_input"],output_names=["x_output", "masks_output"],dynamic_axes=dynamic)torch.onnx.export(model.cpu().eval(),(fake_input.cpu(), fake_mask.cpu()),os.path.join(save_dir, f'static_{i+1}.onnx'),verbose=True,opset_version=17,do_constant_folding=True,input_names=["x_input", "masks_input"],output_names=["x_output", "masks_output"],dynamic_axes=None)

exec="/home/sfy/SFY/camera/TensorRT-8.6.1.6/bin/trtexec"
dir="log"
rm "$dir"/*.txt "$dir"/*.plan: <<'EOF'
# onnx-simplifier 移除冗余节点
for file in "$dir"/*.onnx; doif [[ ! -f "$file" ]]; thenecho "No .onnx files found in $dir."continuefifn=$(basename "$file")python -m onnxsim "$file" "$dir/${fn%.onnx}_simp.onnx"
done
EOFfor file in "$dir"/*.onnx; doif [[ ! -f "$file" ]]; thenecho "No .onnx files found in $dir."continuefifn=$(basename "$file")if [[ $fn == dynamic* ]]; then$exec \--onnx="$file" \--saveEngine="$dir/${fn%.onnx}.plan" \--minShapes=x_input:1x512x768,masks_input:1x768 \--optShapes=x_input:2x512x768,masks_input:2x768 \--maxShapes=x_input:4x512x768,masks_input:4x768 \--verbose \> "$dir/${fn%.onnx}.txt" 2>&1elif [[ $fn == static* ]]; then$exec \--onnx="$file" \--saveEngine="$dir/${fn%.onnx}.plan" \--verbose \> "$dir/${fn%.onnx}.txt" 2>&1elsecontinuefi
done

上述代码对源代码做了简化，仅剥离出造成问题的部分，用于复现和测试 Bug，输入可以看作 1 个 Batch 有 768 张图像序列，每个图像用 512 维特征向量表示；masks 原本是 bool 类型代表 768 张图像是真实数据还是填充数据。

经过测试发现 Bug 由动态模式、F.interpolate、squeeze 组合引发，因此代码推理阶段分为以下 4 种模式：
（1）F.interpolate(mask, size=x.size(-1), mode="nearest") + mask.squeeze(1) 报错
（2）F.interpolate(mask, size=384, mode="nearest") + mask.squeeze(1) 动态模式报错，静态模式通过
（3）mask.squeeze(1) 通过
（4）F.interpolate(mask, size=x.size(-1), mode="nearest") + mask.reshape(b, -1) 或 mask.view(b, -1) 通过

2. 解决方法

用 reshape 或 view 代替 squeeze

3. 原因分析

总结：squeeze 操作需要判断对应维度是否等于 1，而 F.interpolate 改变了张量形状、动态模式引入了维度数值的不确定性，这些使得 onnx 无法确定该维度是否等于 1（即使看上去可以推导出数值），导致添加了 If 节点，而 If 节点在不同状态下分别执行 Squeeze 和 Identity 导致输出形状不统一。采用 Reshape 操作可以直接规避此问题。

下面是不同模式 ONNX 的结构图对比，以及分析细节。

mode1

在这里插入图片描述
mode1 在插值时使用 x.size(-1)，在动态模式下需要通过 Shape 等一系列节点来获取维度信息；在静态模式下 x 所有初始维度固定，把 x.size(-1) 看作固定常量。

mode2

在这里插入图片描述
动态模式下 masks 经过 Resize 维度信息是不确定的，导致需要判断 masks.shape[1] == 1 才能执行 masks.squeeze(1)，问题在于当不等于 1 时执行的是 Identity 导致输出形状不一致。

比较令人费解的是对比 mode1 和 mode2 静态模式，可以发现 Resize 之前的结构完全相同，但是 mode1 在 Resize 之后仍引入了 If 节点导致异常。查到的解释是 squeeze 操作比较保守，mode1 原本 Resize 的维度是动态的依赖 x.size(-1)，即使通过推导将动态转变为静态，但仍保留了动态逻辑（If 节点）。
启用脚本中 onnx-simplifier 移除冗余节点的部分可以去除 mode1 静态模式中的 If 节点，但此方法对动态模式无效。