ONNX Runtime Web如何调用内置Softmax算子获取预测概率

ONNX Runtime Web 已内置Softmax算子实现，无需手动编写JS版Softmax逻辑，也无需修改模型重导出，可直接通过未在公开文档标注的内置算子接口调用，具体实现方式如下：

方案一：推理后直接调用内置Softmax算子（推荐，无模型改动） #

该方案适配你当前已经跑通的推理流程，仅需要在拿到logits输出后追加一步算子调用即可：

• 版本要求：ONNX Runtime Web 版本 >= 1.13.0
• 调用入口：所有内置ONNX标准算子都挂载在全局ort.ops命名空间下，Softmax的入参仅需要输入张量、计算轴两个核心参数
• 参考代码：

// 复用你已有的推理逻辑
const inferenceResults = await bertSession.run(modelInputs);
// 替换为你模型输出的logits对应的字段名，常见值为logits、output_0
const logitsTensor = inferenceResults.logits;
// BERT分类任务默认对最后一维（类别维度）做Softmax计算，axis传-1即可
const probTensor = await ort.ops.softmax({
  input: logitsTensor,
  axis: -1
});
// 取出概率值转为普通JS数组使用
const classProbabilities = await probTensor.getData();

版本兼容提示：1.13及更早的1.x版本中ort.ops下的算子为同步实现，调用时不需要加await，直接接收返回值即可。

方案二：模型导出阶段内置Softmax节点 #

如果不想在前端侧做后处理，可以在导出ONNX模型阶段直接把Softmax加入计算图，推理时直接输出概率值：

• 使用Optimum工具导出Hugging Face BERT分类模型时，指定任务参数为text-classification，导出的模型会自动在输出层追加Softmax节点
• 自定义导出脚本时，在原BERT模型的输出层前追加torch.nn.Softmax(dim=-1)层，再执行ONNX导出即可

补充说明 #

ort.ops下的算子调用入口目前仅在包的TypeScript类型定义文件中有声明，未出现在官方公开教程和API引导文档中。除Softmax外，所有ONNX标准支持的算子都可以通过ort.ops.算子名的格式直接调用，入参规则和ONNX官方算子定义完全一致。

内容的提问来源于stack exchange，提问作者Juan Alberto López Cavallotti