如何在ONNX转换时为Conformer模型正确使用torch._constrain_as_size？

为什么会报错？ #

你遇到的can only constrain range for SymInt错误，是因为直接在forward里用torch._constrain_as_size时，输入维度是普通整数类型（Int），而这个API只接受dynamo追踪时生成的符号整数（SymInt）——只有当维度是动态的、导出时不确定具体值的情况下，才会被转为SymInt。

两种可行的解决方法（无需重新训练模型） #

方法1：导出时直接指定约束（推荐，不用改模型） #

不用修改原模型代码，调用torch.onnx.dynamo_export时，通过constraints参数明确动态维度范围，同时配合dynamic_axes标记动态轴：

import torch
import torchaudio

# 加载已训练好的模型
model = torchaudio.models.ConformerSpeechRecognizer(...)  # 填入你的模型参数
model.load_state_dict(torch.load("trained_model.pth"))
model.eval()

# 构造示例输入，模拟动态序列长度
batch_size = 2
input_dim = 80
seq_len = 300  # 填一个你实际场景范围内的长度即可
inputs = torch.randn(batch_size, input_dim, seq_len)

# 导出ONNX，指定动态轴和约束
with torch.no_grad():
    export_result = torch.onnx.dynamo_export(
        model,
        inputs,
        dynamic_axes={
            "input": {2: "sequence_length"}  # 标记第3维为动态轴
        },
        constraints=[
            "input[2] >= 100",  # 序列长度最小100
            "input[2] <= 500"   # 序列长度最大500
        ]
    )
export_result.save("conformer_onnx_model.onnx")

方法2：用包装类修改forward逻辑（适配已训练模型） #

如果想通过修改模型脚本实现，不需要改动原模型的结构和权重，用包装类对原模型的forward做一层封装，只在dynamo追踪时（维度为SymInt时）添加约束：

class ConstrainedConformer(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.original_model = original_model

    def forward(self, input, lengths=None):
        # 针对输入的序列长度维度（假设shape是[batch, feat, seq_len]）添加约束
        seq_len_dim = input.size(2)
        # 只有当维度是SymInt时才执行约束，避免训练阶段报错
        if isinstance(seq_len_dim, torch.SymInt):
            seq_len_dim = torch._constrain_as_size(seq_len_dim, min=100, max=500)
        # 调用原模型的forward方法
        return self.original_model(input, lengths)

# 包装已训练好的模型
constrained_model = ConstrainedConformer(model)
constrained_model.eval()

# 执行导出
with torch.no_grad():
    export_result = torch.onnx.dynamo_export(
        constrained_model,
        inputs,
        dynamic_axes={"input": {2: "sequence_length"}}
    )
export_result.save("conformer_constrained.onnx")

关键注意点 #

• 约束的必须是动态维度：也就是实际推理时会变化的维度（比如语音序列长度），静态维度（比如特征数、固定的batch size）不需要约束。
• 区分_constrain_as_size和_constrain_as_value：前者用于约束作为张量尺寸的数值（比如序列长度），后者用于约束张量内部的数值。
• 已训练模型完全可以直接适配：不需要重新训练，不管用哪种方法，都只是在导出阶段添加维度约束逻辑，不会改动模型权重。

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