ONNX模型中仅用于控制流的未使用变量推理报错咨询

核心原因 #

PyTorch导出ONNX时会执行**死代码消除（Dead Code Elimination）**优化：任何没有参与最终输出计算链路的输入节点，都会被导出器从ONNX模型中移除。你遇到的报错本质是：导出后的模型里根本不存在x2或布尔标志这类输入节点，ONNX Runtime自然找不到对应的输入名。

解决方法 #

1. 强制让控制输入参与计算链路 #

给控制输入（布尔标志/未使用的输入）添加一个不影响最终结果的无意义运算，让导出器认为该输入是计算链路的必要部分。示例代码：

import torch
from typing import Optional

class SumModule(torch.nn.Module):
    def forward(self, x1, x2, use_x2):
        # 强制让use_x2参与计算，避免被优化掉
        dummy = x1 * 0 + use_x2.float()
        dummy = dummy - dummy  # 运算结果为0，不影响最终输出
        
        if use_x2:
            output = x1 + x2
        else:
            output = x1
        return output

# 导出示例
model = SumModule()
x1_sample = torch.randn(1, 3)
x2_sample = torch.randn(1, 3)
use_x2_sample = torch.tensor(True)

torch.onnx.export(
    model,
    (x1_sample, x2_sample, use_x2_sample),
    "sum_model.onnx",
    input_names=["x1", "x2", "use_x2"],
    output_names=["output"],
    opset_version=14
)

2. 关闭ONNX导出的常量折叠优化 #

通过do_constant_folding=False参数关闭部分死代码消除逻辑，但注意这可能导致模型性能下降，仅建议临时使用：

torch.onnx.export(
    model,
    (x1_sample, x2_sample, use_x2_sample),
    "sum_model.onnx",
    input_names=["x1", "x2", "use_x2"],
    output_names=["output"],
    do_constant_folding=False,
    opset_version=14
)

3. 改用ONNX的Optional输入类型（推荐） #

ONNX 1.10+和PyTorch 1.9+支持Optional[Tensor]类型，可直接替代“输入+布尔标志”的方案，无需额外控制变量：

class SumModule(torch.nn.Module):
    def forward(self, x1, x2: Optional[torch.Tensor] = None):
        if x2 is not None:
            output = x1 + x2
        else:
            output = x1
        return output

# 导出时传入None作为Optional输入的示例
torch.onnx.export(
    model,
    (x1_sample, None),
    "sum_model.onnx",
    input_names=["x1", "x2"],
    output_names=["output"],
    opset_version=14  # 需要opset 14及以上版本
)

推理时：

• 不需要x2时，仅传入x1即可；
• 需要x2时，正常传入该张量。

内容的提问来源于stack exchange，提问作者prozaxx