RKNN千问模板加载模型失败如何解决？

Question

在使用RKNN Toolkit进行大模型部署时，基于“千问”系列模型的转换与加载常出现“Input shape mismatch”错误。该问题通常发生在模型输入维度与RKNN推理要求不一致时，尤其是在Qwen类Transformer模型未正确设置动态轴或输入shape固定不当的情况下。此外，ONNX转RKNN过程中若未合理配置预处理参数（如mean/std、channel_order），也会导致加载失败。此问题表现为调用`rknn.init_runtime()`时报错无法解析输入张量。需结合日志定位具体阶段，检查模型导出脚本中输入shape定义、ONNX优化步骤及RKNN配置参数是否匹配目标硬件平台。

Answer 1

1. 问题背景与现象描述

在使用RKNN Toolkit进行大模型部署时，基于“千问”系列（Qwen）Transformer架构的模型常出现Input shape mismatch错误。该错误通常出现在调用rknn.init_runtime()阶段，表现为无法正确解析输入张量，导致推理初始化失败。

此类问题的根本原因多集中于两个层面：一是模型导出为ONNX格式时输入shape未合理配置，尤其是动态轴（dynamic axes）设置不当；二是从ONNX转换至RKNN过程中预处理参数（如mean/std、channel_order）与目标平台不匹配。

2. 常见错误触发场景

• 静态shape固定错误：将Qwen模型的输入序列长度（如input_ids）设为固定值，但实际推理中变长输入导致维度不一致。
• 动态轴定义缺失：导出ONNX时未声明动态batch_size或sequence_length，造成后续RKNN工具链无法处理可变输入。
• 通道顺序不一致：图像类输入若存在channel_first与channel_last混淆，会引发shape解析异常。
• 预处理配置偏差：RKNN中设置的mean=[128]而实际模型训练使用mean=[0.5]，虽不影响shape数值，但可能导致内部张量重排失败。

3. 分析流程与日志定位方法

当遇到Input shape mismatch时，应按以下步骤逐步排查：

1. 检查ONNX模型输入定义：onnx_model.graph.input
2. 验证ONNX是否包含动态维度（symbolic shape）
3. 查看RKNN转换日志中“Input Nodes”和“Model Input Shape”输出
4. 比对rknn.config()中input_size_list与ONNX输入shape是否对齐
5. 确认rknn.build(do_quantization=True)前的数据预处理设置

4. 模型导出阶段的关键配置示例

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载千问模型
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B")

# 构造动态输入示例
dummy_input = torch.randint(0, 10000, (1, 512))  # batch=1, seq_len=512

# 导出ONNX，关键点：设置动态轴
torch.onnx.export(
    model,
    dummy_input,
    "qwen.onnx",
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "sequence"},
        "logits": {0: "batch", 1: "sequence"}
    },
    opset_version=13
)
    

5. RKNN转换过程中的核心参数配置

参数名	推荐值	说明
mean_values	[[0]]	根据模型训练归一化方式设定
std_values	[[1]]	避免误加标准化导致张量变形
target_platform	rknpu2	确保支持Qwen等大模型
input_size_list	[[1, 512]]	需与ONNX输入shape兼容
quantized_dtype	asymmetric_affine	影响内存布局与shape对齐

6. ONNX优化与兼容性处理建议

部分Qwen模型导出的ONNX存在冗余节点或不兼容算子，建议使用onnx-simplifier进行优化：

pip install onnxsim
python -m onnxsim qwen.onnx qwen_sim.onnx --dynamic-input-shape
    

优化后需重新校验输入shape：

import onnx
onnx_model = onnx.load("qwen_sim.onnx")
print(onnx_model.graph.input[0].type.tensor_type.shape)
    

7. 完整的RKNN构建与运行代码流程

from rknn.api import RKNN

rknn = RKNN(verbose=True)

# 配置参数
rknn.config(
    mean_values=[[0]],
    std_values=[[1]],
    target_platform='rknpu2',
    optimization_level=3
)

# 加载ONNX模型
ret = rknn.load_onnx(model="qwen_sim.onnx")
if ret != 0:
    print("Failed to load ONNX model.")
    exit(ret)

# 构建模型（含量化）
ret = rknn.build(do_quantization=True, dataset_func=get_calibration_data)
if ret != 0:
    print("Failed to build RKNN model.")
    exit(ret)

# 导出并初始化运行时
rknn.export_rknn("qwen.rknn")
ret = rknn.init_runtime()
if ret != 0:
    print("Failed to init runtime, check input shape and pre-process settings.")
    exit(ret)
    

8. 典型错误日志分析与对应解决方案

日志片段	问题根源	解决路径
"Expected shape [1,512], got [1,256]"	输入长度不一致	启用dynamic_axes或padding/truncation
"Channel order not supported"	channel_order配置错误	设置channel_order='normal'
"Quantization failed on node"	预处理参数冲突	关闭量化或调整mean/std
"Cannot find input node 'input_ids'"	节点名称不匹配	检查ONNX输入名与build时绑定关系

9. 使用Mermaid绘制调试流程图

graph TD
    A[开始部署Qwen模型] --> B{ONNX导出成功?}
    B -->|否| C[检查PyTorch导出脚本]
    B -->|是| D[使用onnx-simplifier优化]
    D --> E{是否存在dynamic axes?}
    E -->|否| F[添加动态轴定义]
    E -->|是| G[RKNN配置input_size_list]
    G --> H[调用rknn.build()]
    H --> I{构建失败?}
    I -->|是| J[检查mean/std/channel_order]
    I -->|否| K[调用rknn.init_runtime()]
    K --> L{初始化报错Input shape mismatch?}
    L -->|是| M[比对ONNX输入shape与RKNN配置]
    L -->|否| N[部署成功]
    

10. 跨平台适配与未来优化方向

随着Qwen系列模型向更大参数量演进（如Qwen-72B），在瑞芯微NPU上部署面临更多挑战。除shape匹配外，还需关注分片加载、KV Cache管理、内存带宽瓶颈等问题。建议结合RKNN Toolkit 1.6+版本引入的subgraph partitioning机制，将模型拆分为CPU+NPU协同执行模式，提升整体推理稳定性。

此外，可通过自定义ONNX修改工具动态注入Shape Infer节点，增强工具链对复杂Transformer结构的支持能力，从根本上减少shape推导错误的发生概率。