Pix2Struct HuggingFace模型转ONNX格式维度错误调试求助

问题原因 #

Pix2StructForConditionalGeneration是多模态图文模型，推理时需要同时接收图像输入（pixel_values）和文本输入（input_ids、attention_mask），原代码仅传入了二维的input_ids张量，与模型内部期望的输入维度和数量不匹配，导致导出报错。

修正后的完整代码 #

import torch
import onnx
import onnxruntime
from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor

# 加载模型与处理器
model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-ai2d-base")
processor = Pix2StructProcessor.from_pretrained("google/pix2struct-ai2d-base")
model.eval()

# 准备示例输入：用随机张量模拟图像（实际可替换为PIL格式图像）
dummy_image = torch.randn(1, 3, 224, 224)  # batch_size=1，3通道，224x224尺寸
dummy_text = "Describe this image."  # 示例文本指令

# 通过处理器生成模型所需的全部输入张量
inputs = processor(images=dummy_image, text=dummy_text, return_tensors="pt")

# 导出ONNX模型
onnx_path = "model/pix2struct.onnx"
torch.onnx.export(
    model,
    (inputs["pixel_values"], inputs["input_ids"], inputs["attention_mask"]),
    onnx_path,
    opset_version=16,  # 升级opset版本提升兼容性
    input_names=["pixel_values", "input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "pixel_values": {0: "batch_size", 2: "height", 3: "width"},
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size", 1: "seq_len"}
    }
)

# 验证ONNX模型有效性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)

# ONNX Runtime推理测试
session = onnxruntime.InferenceSession(onnx_path)
output_name = session.get_outputs()[0].name

# 准备推理输入数据（转换为numpy格式）
input_data = {
    "pixel_values": inputs["pixel_values"].numpy(),
    "input_ids": inputs["input_ids"].numpy(),
    "attention_mask": inputs["attention_mask"].numpy()
}

# 执行推理并输出结果
output = session.run([output_name], input_data)
print("推理输出形状:", output[0].shape)
print("模型转换为ONNX成功！")

关键修改说明 #

• 使用Pix2StructProcessor处理输入，确保生成模型要求的全部输入张量，避免维度缺失
• 导出ONNX时传入完整的输入参数元组，而非单一的input_ids张量
• 配置dynamic_axes支持动态维度，让模型适配不同批量大小、图像尺寸和文本序列长度
• 升级opset版本到16，提升算子兼容性与模型功能支持度

内容的提问来源于stack exchange，提问作者Gate computer science learning