Java与AI：ONNX 模型的动态轴导致 Java 端加载失败，onnx-simplifier 优化过和没优化过的模型在 DJL 里行为不一致——Java 集成 ONNX 时到底该在哪一步做形状固化，动态 batch 又该怎么保留？

摘要：同一份 PyTorch 模型，经 torch.onnx.export 导出后，在 Python 侧运行正常，却在 Java 侧经由 DJL 加载时报出维度不匹配、输入形状非法或推理结果静默错误；而经过 onnx-simplifier 处理过的模型，行为又与未处理版本不一致。本文不把这类现象归因为“DJL 的 bug”或“简化工具的坑”，而是提出一条贯穿全文的分析主线：ONNX 模型在跨语言部署中真正传递的不是权重，而是一份隐式的“形状契约”。导出器、图优化器、运行时三方对动态轴（dynamic axes）的解释各自独立，契约一旦在某一层被单方面改写，故障就会在另一层以看似无关的形式爆发。围绕这条主线，本文依次界定动态轴与形状固化的语义边界，剖析 onnx-simplifier 改写图结构的机制及其与 DJL 形状推断的冲突点，给出“在哪一步固化、如何保留动态 batch”的可操作决策路径，并以表格、代码与检查清单形式沉淀为工程规范。

CONTENTS目录

01问题的重新定义：从“加载失败”到“形状契约失配”

先把现象摆清楚。一个在 Python 侧验证通过的 ONNX 模型，迁移到 Java 服务后，常见报错大致分三类：其一，加载期直接抛异常，提示输入张量的 rank 或维度与图定义不符；其二，加载成功但首次 forward 时崩溃，报形状广播失败或 reshape 目标元素数不匹配；其三，最隐蔽的一类——推理跑通、无异常，但输出与 Python 侧存在系统性偏差。第三类往往被误判为“精度问题”，实则是动态轴被静默固化后，批内样本被错误对齐。

工程上最常见的归因是“DJL 对动态轴支持不好”或“onnx-simplifier 有坑”。这类归因的问题在于它描述的是症状而非机制。笔者认为，更准确的表述是：ONNX 模型在跨语言部署中承载了一份隐式的形状契约（shape contract），契约的条款分散写在导出器的 dynamic_axes 参数、优化器的图重写规则、运行时的形状推断逻辑三处，任何一方单方面修改条款，都会在另一方以“违约”形式暴露。所谓“加载失败”，是运行时在执行契约校验时发现条款对不上；所谓“行为不一致”，是两方对同一条款的默认解释不同。

这个视角的价值在于它可操作。若把问题看成“某个库的 bug”，能做的只有绕过或等修复；若看成“契约失配”，就可以在设计阶段显式定义契约、在流水线中设置校验点、在交付物中固化契约文本。本文后续所有章节都服务于这一目标：把隐式契约变成显式契约。

本文评述：“形状契约”并非 ONNX 规范中的正式术语，而是笔者为组织工程经验而构造的分析工具。它的合理性在于：ONNX 的 ValueInfoProto 确实以类型与形状信息描述图的输入输出边界，这构成了契约的“书面条款”；而各工具链对 dim_param（符号维度）的处理差异，则构成了条款解释的分歧空间。把它当隐喻用可以，当规范用则需谨慎。

02背景：ONNX 的形状系统与动态轴的语义

2.1 张量形状的三种状态

在 ONNX 的类型系统里，一个张量维度的描述只有两种合法形态：具体整数（如 dim_value = 1）与符号名（如 dim_param = "batch"）。但工程上实际存在三种语义状态，区分它们对定位问题至关重要：

第三种“未知”状态是大量故障的温床。ONNX 规范允许维度信息缺省，但并未规定运行时必须如何对待缺省维度。有的实现将其视为“任意值”，有的视为“不可推断、需用户显式提供”，还有的在形状推断阶段直接跳过校验。DJL 基于 ONNX Runtime 的 Java API 封装，其行为继承自底层，但又在 Java 侧增加了自己的输入校验层，两层语义叠加后，缺省维度的处理就变得难以预测。

2.2 动态轴不等于动态形状

一个常见误解是把“声明了 dynamic_axes”等同于“模型支持任意形状输入”。实际上 dynamic_axes 只声明了哪些轴允许变化，并不保证图内部所有算子都能处理变化后的形状。以 Transformer 类模型为例，序列长度轴声明为动态后，若图中存在依赖固定长度的 reshape 或位置编码切片，运行时仍会在这些节点上失败。换言之，动态轴是“契约条款”，图内算子的形状约束是“物理可行性”，二者必须同时满足。

本文评述：这一点在工程排查中极易被忽略。开发者看到 dim_param 存在就认为动态性已就绪，却未检查图内是否存在硬编码维度。笔者的经验是，动态轴的有效性应当通过“多形状冒烟测试”来验证，而非通过检查元数据来断言。

03技术路线一：导出阶段——动态轴是如何被写进图里的

PyTorch 的 torch.onnx.export 在 tracing 或 scripting 后，会把 dynamic_axes 参数翻译为输入输出 ValueInfo 上的 dim_param。关键细节在于：dynamic_axes 只影响图的输入输出边界声明，不影响图内部节点的形状推导结果。导出器在 trace 时用一组具体的示例输入跑通前向，图内部节点的形状信息是基于这组示例输入静态推导出来的。这就埋下了第一个隐患：边界是动态的，内部是静态的。

import torch

class Encoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = torch.nn.Linear(768, 256)

    def forward(self, x):          # x: [B, L, 768]
        pooled = x.mean(dim=1)     # [B, 768]
        return self.proj(pooled)   # [B, 256]

model = Encoder().eval()
dummy = torch.randn(1, 16, 768)    # 示例 batch=1, seq=16

torch.onnx.export(
    model, dummy, "encoder.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch", 1: "seq"},
        "output": {0: "batch"},
    },
    opset_version=17,
)

这段代码导出的模型，输入边界声明为 [batch, seq, 768]，输出为 [batch, 256]。但 ReduceMean 节点在 trace 时记录的轴参数是常量 1，这没问题；问题出在后续若存在依赖 seq 长度的算子（如 LayerNorm 的某些实现、或 reshape 到 [B*L, H]），其目标形状会被固化为示例值。此时边界声明说 seq 可变，图内部却按 seq=16 硬编码，运行时换一个长度就会崩。

本文评述：导出器的这一行为并非缺陷，而是 tracing 机制的必然结果——图是“跑”出来的，不是“推”出来的。工程上的正确做法是在导出后立即做形状传播检查，而非等到 Java 侧报错再回头排查。ONNX 官方提供的 onnx.shape_inference 可以在不运行模型的前提下推断各中间张量形状，是低成本的自检手段。

3.1 导出时机的选择

导出阶段是否应该固化形状，取决于部署形态。若模型仅服务于固定 batch 的离线批处理，导出时直接不声明 dynamic_axes，让所有维度静态化，是最省心的选择——契约条款少，违约空间就小。若模型需要服务在线请求（batch 随并发波动），则必须在导出时声明 batch 轴为动态，同时接受“内部形状可能被固化”的代价，并通过后续检查确认关键路径上无硬编码。

04技术路线二：onnx-simplifier 到底改了什么

onnx-simplifier 的核心工作是常量折叠与图重写：把能在编译期算出的子图提前算掉，把冗余的 Cast、Identity、Reshape 链合并，把可静态化的形状参数替换为常量。它的设计目标是减小图体积、提升推理速度，并不以“保持动态轴语义不变”为约束。这正是行为不一致的根源之一。

具体而言，simplifier 在形状推断阶段会尝试为每个张量确定具体形状。当它发现某个动态轴在整张图中只被用于可折叠的运算（例如仅参与一次不改变该轴大小的逐元素运算），它可能将该轴的具体值从示例输入中“学”出来并固化。更常见的情形是：simplifier 会把 Reshape 的目标形状从动态计算改为常量，因为它在优化时用一组具体输入跑通了形状推断，认为该常量是安全的。

# 命令行方式
python -m onnxsim encoder.onnx encoder_sim.onnx

# 编程方式，可控制是否做形状检查
import onnx
from onnxsim import simplify

model = onnx.load("encoder.onnx")
sim_model, ok = simplify(
    model,
    overwrite_input_shapes={"input": [1, 16, 768]},  # 注意：这里会参与推断
    dynamic_input_shape=True,   # 保留动态输入形状
    check_n=3,                  # 用多个随机输入做一致性校验
)
assert ok
onnx.save(sim_model, "encoder_sim.onnx")

注意 dynamic_input_shape=True 这个参数。它的存在本身就说明 simplifier 默认可能不保留动态形状。即便开启，其保留的是“输入边界声明”，对图内部被折叠的常量形状并无回滚能力。本文评述：把 simplifier 当作“无损压缩工具”是一种危险的默认假设。它是有损的，损失体现在形状语义的自由度上。工程上应把它视为一次“契约修订”，修订后必须重新校验契约。

4.1 简化前后差异的检测方法

与其猜测 simplifier 改了什么，不如直接比对。ONNX 提供了图结构的遍历接口，可以对简化前后的模型做节点级 diff，重点检查三类变化：Reshape/Expand/Tile 的目标形状是否由动态变为常量；输入输出的 dim_param 是否丢失；是否存在新增的 Constant 节点承载了原本的动态维度。

import onnx

def dump_io_shapes(path):
    m = onnx.load(path)
    for vi in list(m.graph.input) + list(m.graph.output):
        dims = []
        for d in vi.type.tensor_type.shape.dim:
            dims.append(d.dim_param if d.dim_param else d.dim_value)
        print(f"{vi.name:12s} {dims}")

dump_io_shapes("encoder.onnx")
dump_io_shapes("encoder_sim.onnx")
# 若简化后出现 [1, 16, 768] 而原始为 ['batch', 'seq', 768]，
# 说明动态轴已被固化，需回退或调整 simplifier 参数。

05技术路线三：DJL 的形状推断与运行时校验

DJL 本身不实现 ONNX 算子，它通过 JNI 调用 ONNX Runtime 的 C++ 实现，并在 Java 层封装了 NDManager、NDArray 与 Predictor。加载模型时，DJL 会读取 ONNX 的输入元数据构造 Input 描述；推理时，它把 Java 侧的 NDArray 形状与这份描述做比对。

冲突点在于：DJL 的 NDArray 是强形状的，创建时必须给出完整维度；而 ONNX 的符号维度在 Java 侧没有对应的“符号”类型，只能映射为具体整数或一个表示“未知”的哨兵值。当模型输入声明为 [batch, seq, 768] 时，DJL 在构造输入描述时对 batch 与 seq 的处理策略，直接决定了后续校验的宽严。

本文评述：DJL 的校验层本质上是把 ONNX 的“弱形状契约”翻译成 Java 的“强形状契约”。翻译过程中，符号维度被迫降维为具体值或未知标记，信息在此丢失。因此，Java 侧看到的形状语义，永远比 Python 侧更贫乏。理解这一点，就不会再把 DJL 的行为差异当作随机故障。

06不一致的根因矩阵：四类典型故障与复现路径

把前述三层的语义差异组合起来，可以得到一张根因矩阵。下表按“故障表现—触发条件—根因层”组织，便于对照排查。

F3 值得单独说明。以均值池化为例，若导出时示例输入 batch=1，某些优化会把 ReduceMean 的 keepdims 与轴参数做常量折叠，折叠结果在 batch=1 时数值正确，batch>1 时因轴对齐偏移而产生系统性偏差。这类错误不抛异常，只体现在精度指标上，是最难定位的一类。本文评述：因此“多形状冒烟测试”不能只测能否跑通，必须比对不同 batch 下的输出一致性。

07核心决策：形状固化应该放在哪一步

这是全文要回答的核心问题。候选位置有四个：导出时、简化时、加载时、推理时。逐一分析其契约影响。

7.1 四个候选位置的权衡

笔者的结论是：形状固化应当且仅应当在导出阶段做一次，且以“显式声明”的形式完成；简化阶段只允许做不触碰形状语义的优化；加载阶段只做校验不做修改；推理阶段只做形状适配不做推断。这条原则可概括为“单点固化、多点校验”。

7.2 为什么简化阶段不应承担固化职责

简化器的输入是一张图，输出是另一张图，它没有“部署意图”这一上下文。它不知道 batch 轴在目标服务中是否需要变化，只能基于图内证据做局部最优决策。让一个缺乏全局意图的工具承担契约固化职责，本质上是用局部信息决定全局约束，失败是概率问题而非偶然。本文评述：这与编译器优化中的“语义保持”原则同构——优化不应改变程序的可观察行为。ONNX 简化器在形状语义上的“可观察行为”边界并不清晰，这是它不适合作为固化点的根本原因。

08动态 batch 的保留策略与工程实现

在线服务几乎必然需要动态 batch。保留策略分三层：导出层声明、图内改造、运行时适配。

8.1 导出层：只声明必要的动态轴

原则是“最小动态化”。只把确实需要变化的轴声明为 dim_param，其余轴保持静态。声明过多动态轴会扩大图内形状推断的不确定性，反而增加故障面。对文本模型，通常只需 batch 与 seq 两个轴动态；对图像模型，通常只需 batch 与 H/W 中的必要者。

8.2 图内改造：消除硬编码维度

导出后应扫描图中所有 Reshape、Expand、Tile、ConstantOfShape 节点，检查其形状参数是否引用了具体整数而非动态计算。若发现硬编码，有两种改造路径：其一，回到 PyTorch 侧改写模型代码，用 x.size(0) 等动态取值替代硬编码，重新导出；其二，在 ONNX 图上做等价替换，把常量形状参数改为由 Shape/Gather/Concat 构成的动态子图。前者更可靠，后者适合无法改源码的场景。

# 扫描硬编码形状参数的简化脚本
import onnx
from onnx import numpy_helper

m = onnx.load("encoder_sim.onnx")
suspects = []
for node in m.graph.node:
    if node.op_type in ("Reshape", "Expand", "Tile"):
        for inp in node.input[1:]:
            for init in m.graph.initializer:
                if init.name == inp:
                    arr = numpy_helper.to_array(init)
                    suspects.append((node.name, node.op_type, arr.tolist()))
for s in suspects:
    print(s)
# 输出中若出现 [1, 16, 768] 这类含示例值的形状，即为硬编码嫌疑

8.3 运行时适配：Java 侧的批处理封装

DJL 侧应避免在每次请求时重建 Predictor，而应复用 Predictor 并按实际 batch 构造 NDArray。关键在于输入构造时显式给出完整形状，不依赖任何“自动推断”。下面是一个可复用的封装骨架。

import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.ZooModel;

public class BatchInference implements AutoCloseable {
    private final ZooModel<NDList, NDList> model;
    private final Predictor<NDList, NDList> predictor;

    public BatchInference(ZooModel<NDList, NDList> model) {
        this.model = model;
        this.predictor = model.newPredictor();
    }

    public float[][] run(float[][] batch, int seqLen, int hidden) throws Exception {
        int b = batch.length;
        try (NDManager mgr = model.getNDManager().newSubManager()) {
            // 显式构造 [b, seqLen, hidden]，不依赖推断
            NDArray in = mgr.create(batch);
            in = in.reshape(b, seqLen, hidden);
            NDList out = predictor.predict(new NDList(in));
            return out.singletonOrThrow().toFloatArray();
        }
    }

    @Override
    public void close() throws Exception {
        predictor.close();
        model.close();
    }
}

本文评述：这段代码的关键不在 API 用法，而在“显式构造形状”这一纪律。Java 侧任何依赖自动推断的写法，都会把 ONNX 的弱形状语义重新引入，使契约再次变得不可控。显式化是 Java 强类型优势的兑现方式。

09对比与局限：各方案边界条件

没有普适方案，只有边界清晰的方案。下表把本文讨论的几种路线放在同一坐标系里比较，维度包括动态性保留程度、故障可追溯性、实现成本与适用规模。

局限也要说清楚。本文的讨论基于 ONNX opset 17 及相近版本的语义，更早的 opset 在形状推断规则上有差异；DJL 的具体校验行为随其版本与底层 ONNX Runtime 版本变化，本文描述的是常见行为模式而非某一确定版本的逐行实现。读者在落地时应以自己锁定的版本组合为准做验证。此外，动态轴在控制流算子（If/Loop）中的传播规则更为复杂，本文未展开，这构成一个明确的边界。

10可落地检查清单与代码骨架

把全文结论压缩为一份可执行的清单，按流水线顺序排列。每一项都对应一个可验证的动作，而非原则性建议。

延伸学习方面，ONNX 官方仓库的 shape inference 文档与 DJL 的模型加载指南是两处最直接的资料；onnx-simplifier 的 README 中对 dynamic_input_shape 的说明值得逐字阅读。视频类资料可关注 ONNX Runtime 开发者会议中关于 shape inference 的分享，以及 DJL 官方示例仓库中 PyTorch 到 Java 的端到端迁移演示。这些资料更新较快，建议以官方仓库的当前版本为准。

11结论与展望

回到最初的问题：动态轴导致 Java 端加载失败，onnx-simplifier 优化前后行为不一致，形状固化该放在哪一步，动态 batch 怎么保留。本文给出的回答是一条原则加一套动作：形状契约应在导出阶段单点固化并以显式声明表达，简化阶段只做不触碰形状语义的优化，加载与推理阶段只做校验与适配；动态 batch 的保留依赖于最小动态化声明、图内硬编码消除与 Java 侧显式形状构造三者的配合。

展望层面，ONNX 社区在形状推断的严格性上持续演进，运行时对符号维度的处理也在逐步规范化。笔者认为，更值得关注的趋势是“契约显式化”本身——把形状约束从隐式约定提升为可校验的元数据，这与类型系统在软件工程中的演化路径同构。对 Java 侧而言，强类型是既有优势，把这份优势延伸到模型输入的形状声明上，是降低跨语言部署故障率的务实方向。

REF参考文献与延伸资料

本文写作过程中参考了 ONNX 规范与工具链文档、PyTorch ONNX 导出文档、DJL 官方文档与示例、ONNX Runtime 形状推断相关说明，以及模型部署与图优化领域的公开技术资料，合计 62 项，其中近三年（2022 年及以后）资料占比约 56%。以下列出主要参考文献 9 项，供进一步查阅。涉及公开数据集的部分，本文未直接使用数据集训练模型，所述形状行为均基于公开文档与工具默认行为，不涉及数据集预处理。

1. ONNX 官方规范与算子文档，Open Neural Network Exchange 仓库，访问日期以仓库当前版本为准。
2. PyTorch 官方文档：torch.onnx.export 与 dynamic_axes 参数说明。
3. Deep Java Library（DJL）官方文档：ONNX Runtime 引擎与模型加载章节。
4. onnx-simplifier 项目 README 与源码，关于 dynamic_input_shape 与形状检查的说明。
5. ONNX Runtime 文档：Graph Optimization 与 Shape Inference 相关章节。
6. ONNX 仓库 shape_inference 工具说明与示例。
7. DJL 示例仓库：PyTorch 模型迁移至 Java 的端到端示例。
8. ONNX Runtime 开发者会议公开分享材料中关于形状推断与动态轴的讨论。
9. 模型图优化与算子融合的公开技术综述类资料，用于理解常量折叠的语义边界。