HuggingFace ALBERT问答模型导出ONNX后推理维度不匹配问题

尝试将albert-xlarge-v2-squad-v2模型导出为ONNX格式后，仅能使用导出时相同长度的问题和上下文进行推理，更换不同长度的输入就会报错：

InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: attention_mask for the following indices
index: 1 Got: 15 Expected: 26
Please fix either the inputs or the model.

原因 #

直接用固定长度的输入张量导出ONNX时，PyTorch会将输入的维度（如序列长度）固化到模型中，导致后续无法适配不同长度的输入。

解决方法 #

导出ONNX时通过dynamic_axes参数声明动态维度，允许输入的序列长度可变。修改后的导出代码如下：

import torch.onnx
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

tokenizer = AutoTokenizer.from_pretrained("ktrapeznikov/albert-xlarge-v2-squad-v2")
model = AutoModelForQuestionAnswering.from_pretrained("ktrapeznikov/albert-xlarge-v2-squad-v2")

# 生成任意输入示例（仅用于导出时的形状参考）
question = "Where do I live?"
context = "My name is Peter from Australia"
inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt")

input_names = ["input_ids", "attention_mask"]
output_names = ["start_scores", "end_scores"]

# 配置动态维度：第0维为batch size，第1维为序列长度
dynamic_axes = {
    "input_ids": {0: "batch_size", 1: "seq_len"},
    "attention_mask": {0: "batch_size", 1: "seq_len"},
    "start_scores": {0: "batch_size", 1: "seq_len"},
    "end_scores": {0: "batch_size", 1: "seq_len"}
}

torch.onnx.export(
    model,
    (inputs['input_ids'], inputs['attention_mask']),
    "albert_qa_dynamic.onnx",
    input_names=input_names,
    output_names=output_names,
    opset_version=11,
    dynamic_axes=dynamic_axes  # 关键：启用动态维度
)

print("支持动态序列长度的ONNX模型已保存。")

验证推理 #

使用修改后的ONNX模型，即可正常处理任意长度的问题和上下文，推理代码无需改动：

import onnxruntime
import numpy as np
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ktrapeznikov/albert-xlarge-v2-squad-v2")
ort_session = onnxruntime.InferenceSession("albert_qa_dynamic.onnx")

# 测试不同长度的输入
question = "What's my name?"
context = "I'm Alice, and I work as a software engineer in Paris."

inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt")
ort_inputs = {
    ort_session.get_inputs()[0].name: inputs['input_ids'].cpu().numpy(),
    ort_session.get_inputs()[1].name: inputs['attention_mask'].cpu().numpy()
}
ort_outputs = ort_session.run(None, ort_inputs)

start_index = ort_outputs[0].argmax()
end_index = ort_outputs[1].argmax() + 1
answer_tokens = inputs['input_ids'][0][start_index:end_index]
answer = tokenizer.decode(answer_tokens)

print(answer)  # 输出：alice

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