📚 课程目录第4章输入输出张量形状与数据类型

4、输入输出张量形状与数据类型：shape、dtype的定义、动态维度与静态维度、类型映射（float32、int64等）

在ONNX模型的部署与推理过程中，输入输出张量的形状（shape）与数据类型（dtype）是模型描述文件（.onnx）中最重要的元信息之一。它们不仅决定了推理引擎如何分配内存、执行算子，还直接关系到前后处理代码的编写、动态尺寸的适配以及跨框架的兼容性。本章将系统梳理shape与dtype的定义方式、动态与静态维度的区别、ONNX与主流框架（PyTorch、TensorFlow）之间的类型映射关系，并结合实际部署场景给出排查与修正方法。

上图展示了本章的核心知识框架：形状定义、维度类型与数据类型映射三者共同构成了ONNX模型输入输出描述的基础。下面我们将逐一深入解析。

4.1 形状（shape）的定义与结构

核心原理：在ONNX的ValueInfoProto消息中，每个输入或输出张量都通过type.tensor_type.shape字段来描述其形状。形状由一组维度（dimension）组成，每个维度可以是固定的数值（dim_value），也可以是动态的符号名称（dim_param）。例如，一个典型的图像分类模型输入形状为[1, 3, 224, 224]，表示批量大小为1、通道数为3、高度和宽度均为224像素。而一个支持可变批处理大小的模型，其输入形状可能被定义为["batch", 3, 224, 224]，其中"batch"是一个符号维度。

判断方法：在部署前，应使用onnx.shape_inference工具或直接解析模型文件来查看输入输出的实际形状定义。推荐使用Python脚本快速检查：

import onnx
model = onnx.load("model.onnx")
for inp in model.graph.input:
    print("Input:", inp.name)
    shape = inp.type.tensor_type.shape
    for dim in shape.dim:
        if dim.HasField("dim_value"):
            print("  dim_value:", dim.dim_value)
        elif dim.HasField("dim_param"):
            print("  dim_param:", dim.dim_param)
        else:
            print("  unknown dim")

案例拆解：某目标检测模型在导出时使用了固定形状[1, 3, 640, 640]，但在实际生产环境中需要处理不同分辨率的输入图像。部署到ONNX Runtime后，当输入尺寸变为[1, 3, 800, 600]时，推理直接报错Invalid shape for input。排查过程如下：

首先检查模型输入定义，确认shape为固定值；
使用onnxruntime.tools.make_dynamic_shape_fixed或手动修改模型，将高度和宽度维度改为dim_param；
重新导出或使用onnx.shape_inference验证修改后的模型；
在推理代码中动态传入实际尺寸，并确保预处理（resize、pad）与模型期望一致。

修改后的模型输入定义为["batch", 3, "H", "W"]，其中H和W为动态维度，推理时通过ort_session.run(None, {"input": input_tensor})传入任意尺寸的张量即可。

风险提示：动态维度虽然灵活，但会降低部分算子的优化效率（如卷积的im2col、矩阵乘法的分块策略）。在性能敏感场景下，建议将动态维度限制在最小必要范围内，例如仅batch维度动态，而H/W保持固定。

4.2 动态维度与静态维度的工程取舍

核心原理：静态维度（static dimension）在模型编译时即确定，推理引擎可以提前分配固定大小的内存、应用算子融合和内存复用优化。动态维度（dynamic dimension）则允许在运行时指定具体数值，但会引入额外的形状推导开销，并可能阻止某些图优化。ONNX规范中，维度值为0通常表示该维度是动态的（即任意大小），而使用dim_param可以赋予动态维度一个语义化名称。

操作流程：在将PyTorch或TensorFlow模型导出为ONNX时，可以通过dynamic_axes参数控制哪些维度是动态的。例如：

import torch
# 导出时指定batch和序列长度为动态
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "seq_len"},
        "attention_mask": {0: "batch", 1: "seq_len"},
        "logits": {0: "batch", 1: "seq_len"}
    }
)

案例拆解：某自然语言处理模型在导出时未设置动态维度，导致部署后无法处理超过固定长度（如128）的文本序列。通过分析线上日志发现，当输入序列长度超过128时，推理引擎抛出Shape mismatch错误。解决步骤：

重新导出模型，将序列长度维度设为动态；
在预处理阶段对输入文本进行padding或truncation，使其长度与模型动态范围匹配；
在推理服务中增加输入长度校验，避免极端长度导致内存溢出。

修改后，模型可处理长度在1~512之间的任意序列，但推理延迟随长度线性增长。通过压测发现，动态维度带来的性能损失约为5%~8%，在可接受范围内。

维度类型	优点	缺点	适用场景
静态维度	推理速度快、内存占用稳定、算子优化充分	灵活性差、需固定输入尺寸	图像分类、固定分辨率检测、批量推理
动态维度	适应多种输入尺寸、无需重新导出	性能略降、内存分配频繁、部分算子不支持	NLP变长序列、目标检测多尺度、视频流处理
混合维度	兼顾灵活性与性能	配置复杂、需仔细设计	batch动态+空间尺寸固定、或反之

经验建议：在模型导出前，先明确部署场景的输入范围。如果输入尺寸变化不大，优先使用静态维度；如果必须支持多种尺寸，建议将最频繁使用的尺寸设为静态，其他尺寸通过动态维度兜底。

4.3 数据类型（dtype）与ONNX类型映射

核心原理：ONNX定义了丰富的数据类型，包括浮点型（float16、float32、float64）、整型（int8、int16、int32、int64）、无符号整型（uint8、uint16、uint32、uint64）、布尔型（bool）以及字符串型（string）。在模型描述中，每个张量通过tensor_type.elem_type字段指定数据类型，其值为枚举编号（如1表示float32，7表示int64）。

类型映射关系：不同深度学习框架与ONNX之间的类型对应关系如下表所示：

ONNX类型	枚举值	PyTorch	TensorFlow	NumPy
float32	1	torch.float32	tf.float32	np.float32
float16	10	torch.float16	tf.float16	np.float16
float64	11	torch.float64	tf.float64	np.float64
int8	3	torch.int8	tf.int8	np.int8
int16	5	torch.int16	tf.int16	np.int16
int32	6	torch.int32	tf.int32	np.int32
int64	7	torch.int64	tf.int64	np.int64
uint8	2	torch.uint8	tf.uint8	np.uint8
bool	9	torch.bool	tf.bool	np.bool_

判断方法：在推理代码中，必须确保输入张量的dtype与模型定义一致。例如，ONNX Runtime对dtype不匹配会直接抛出Type mismatch错误。常见的错误包括：PyTorch默认使用float32，但模型期望int64（如索引输入）；或NumPy数组默认为int64，而模型期望int32。

案例拆解：某推荐系统模型使用PyTorch导出，输入包括用户ID（int64）和物品ID（int64）。部署时使用Java后端调用ONNX Runtime，Java端构造的long[]数组被转换为int64张量，推理正常。但后续优化时，将用户ID改为int32以节省内存，结果推理失败。排查日志显示：

Error: Unexpected input data type. Actual: (tensor(int32)) , expected: (tensor(int64)).

解决方法是修改模型输入类型为int32，或在Java端将int32数组强制转换为int64后再传入。最终选择修改模型，因为int32足以覆盖用户ID范围，且内存占用减半。

关键要点：类型映射不仅仅是数值转换，还涉及精度损失风险。例如，将float32转换为float16可能导致精度下降，在目标检测等任务中可能引发边界框回归偏差。建议在转换后使用验证集评估模型精度，确保在可接受范围内。

4.4 常见误区与排查清单

核心原理：shape与dtype问题在模型部署中极为常见，且往往在模型转换阶段不易察觉，直到推理阶段才暴露。以下总结了几类高频问题及其排查思路。

常见误区：

误区一：认为动态维度越多越好。实际上，过多的动态维度会显著降低推理性能，且部分自定义算子不支持动态形状。
误区二：忽略dtype的隐式转换。某些框架在导出时会自动将float64降为float32，但若模型内部有精度敏感算子（如Softmax），可能导致结果偏差。
误区三：在预处理阶段使用与模型训练时不同的数据类型。例如，训练时图像归一化使用float32，但部署时使用uint8直接输入，导致输出错误。

排查流程：当推理出现shape或dtype相关错误时，可按以下步骤快速定位：

使用onnxruntime.get_providers()确认当前使用的执行提供程序（CPU/CUDA/TensorRT）；
打印模型输入输出的实际shape与dtype：ort_session.get_inputs()[0].shape、ort_session.get_inputs()[0].type；
对比实际传入张量的shape与dtype，使用np.ndarray.shape和np.ndarray.dtype检查；
若shape不匹配，检查预处理代码中的resize、crop、padding逻辑；
若dtype不匹配，使用tensor.astype(np.float32)或torch.tensor(..., dtype=torch.float32)进行显式转换。

案例复盘：某自动驾驶感知模型在GPU上部署时，输入图像为[1, 3, 1280, 720]，dtype为float32。但实际推理时，图像从相机采集后为uint8格式，未做转换直接输入，导致输出检测框全部偏移。通过添加image = image.astype(np.float32) / 255.0归一化步骤后，问题解决。复盘发现，训练时数据加载器已包含归一化，但部署代码遗漏了该步骤。

风险提示：在模型转换过程中，某些算子（如Resize、Cast）可能会隐式改变张量的shape或dtype。建议在导出后使用onnx.checker.check_model进行完整性校验，并使用onnxruntime的InferenceSession进行冒烟测试，确保模型在目标环境中可正常推理。

综上所述，shape与dtype是ONNX模型部署的基石。掌握静态与动态维度的定义方式、熟悉类型映射关系、建立系统的排查流程，能够显著降低部署过程中的试错成本，提升模型上线的效率与稳定性。