PyTorch转ONNX遇异常：动态轴配置错误还是需Torch Script？

想要将以下Conformer语音识别模型转换为ONNX格式：

import torch
import torchaudio

class ConformerSpeechRecognizer(torch.nn.Module):
    def __init__(self,
                 kernel_size,
                 ffn_dim: int,
                 feature_vector_size: int,
                 hidden_layer_size: int,
                 num_layers: int,
                 num_heads: int,
                 dropout: float,
                 depthwise_conv_kernel_size: int,
                 vocabulary_size: int):
        super().__init__()
        self.vocabulary_size = vocabulary_size
        self.cnn_ = torch.nn.Sequential(
            torch.nn.BatchNorm1d(num_features=feature_vector_size),
            torch.nn.Conv1d(
                in_channels=feature_vector_size,
                out_channels=hidden_layer_size,
                bias=False,
                kernel_size=(kernel_size,),
                padding='same'
            ),
            torch.nn.BatchNorm1d(num_features=hidden_layer_size)
        )
        self.conformer_ = torchaudio.models.Conformer(
            input_dim=hidden_layer_size,
            num_heads=num_heads,
            num_layers=num_layers,
            ffn_dim=ffn_dim,
            depthwise_conv_kernel_size=depthwise_conv_kernel_size,
            dropout=dropout
        )
        self.proba_appoximator_ = torch.nn.Linear(
            in_features=hidden_layer_size,
            out_features=vocabulary_size
        )
    def forward(self, inputs: torch.Tensor, input_lenghts: torch.Tensor) -> torch.Tensor:
        hidden_states = torch.nn.functional.gelu(self.cnn_(inputs.permute(0, 2, 1)))
        hidden_states, _ = self.conformer_(hidden_states.permute(0, 2, 1), input_lenghts)
        output_logits = self.proba_appoximator_(hidden_states)
        return output_logits

转换时使用的导出代码如下：

torch.onnx.export(model.cpu(),
                  f="../model.onnx",
                  input_names=["inputs", "input_lenghts"],
                  output_names=["logits"],
                  args=({ "inputs": torch.ones(1, 1, 13), "input_lenghts": torch.ones(1, dtype=torch.long) }),
                  dynamic_axes={
                    "inputs": { 0: "batch_count", 1: "batch_item_length" },
                    "input_lenghts": { 0: "batch_count" },
                    "logits": { 0: "batch_count", 1: "batch_item_length" }
                  },
                  verbose=True
                )

遇到的问题：

• 不使用torch.jit.script(model)时，转换仅出现padding警告，但测试模型会触发Reshape节点运行时异常；修改输入batch维度正常，但修改序列长度（第二个维度）时转换失败，不符合动态轴预期。
• 先将模型转为Torch Script再导出，会触发ScalarType UNKNOWN_SCALAR运行时错误。

1. 核心错误：args参数格式错误 #

torch.onnx.export的args参数要求传入模型forward方法的位置参数元组，而非关键字参数字典。你当前用字典传递输入，会导致ONNX导出器无法正确绑定输入张量与动态轴的对应关系，这是序列长度维度动态性失效的主要原因。

2. 示例输入序列长度过小 #

你使用的示例输入序列长度为1（torch.ones(1,1,13)的第二个维度），这种极端小的序列长度会导致Conformer内部部分层（如深度可分离卷积）的计算逻辑在导出时被固化为固定形状，无法适配动态序列长度。建议使用更贴近实际场景的序列长度，比如32。

3. Torch Script导出的类型推断问题 #

直接脚本化整个模型时，torchaudio内置的Conformer部分操作可能存在类型推断不明确的情况，引发UNKNOWN_SCALAR错误。无需强制使用Torch Script，调整导出参数即可解决问题。

修改后的导出代码 #

# 先实例化模型（假设参数已正确设置）
model = ConformerSpeechRecognizer(
    kernel_size=3,
    ffn_dim=256,
    feature_vector_size=13,
    hidden_layer_size=64,
    num_layers=4,
    num_heads=2,
    dropout=0.1,
    depthwise_conv_kernel_size=31,
    vocabulary_size=1000
)
model.eval()  # 必须切换到eval模式，避免dropout等训练层影响

# 构造更合理的示例输入：batch=1，seq_len=32，feature_dim=13
dummy_input = torch.ones(1, 32, 13)
dummy_lengths = torch.tensor([32], dtype=torch.long)

torch.onnx.export(
    model.cpu(),
    f="../model.onnx",
    args=(dummy_input, dummy_lengths),  # 使用位置参数元组
    input_names=["inputs", "input_lenghts"],
    output_names=["logits"],
    dynamic_axes={
        "inputs": {0: "batch_count", 1: "batch_item_length"},
        "input_lenghts": {0: "batch_count"},
        "logits": {0: "batch_count", 1: "batch_item_length"}
    },
    opset_version=17,  # 使用更高版本opset，支持更多动态形状操作
    do_constant_folding=False,  # 禁用常量折叠，避免动态形状被固化
    verbose=True
)

额外注意事项 #

• 导出前必须将模型切换到eval()模式，避免训练时的随机层（如dropout）干扰导出逻辑。
• 确保input_lenghts的数值与输入序列的实际长度一致（示例中是32），否则Conformer内部的掩码计算会出现异常。
• 若仍存在Reshape报错，可以检查ONNX模型的节点结构（用Netron工具可视化），确认是否有未被正确标记为动态的维度。

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