如何最有效地将一个 ONNX 模型的输出提供给 Java 中的另一个？

问题描述

我有两个使用 ONNX 构建的模型，模型 A 和模型 B。我使用 ONNX Runtime Java API 加载这些模型并使用它们进行推理。工作流程是我需要用模型 A 计算预测，然后将模型 A 的结果输入模型 B：

x -> A(x) -> B(A(x)) -> y

当我调用 resultFromA = A.run(inputs)(OrtSession.run) 时，API 返回 Result。 理想情况下，我想获取该结果并调用 B.run(resultFromA)，但 run 只接受 Map<String,OnnxTensor> inputs。我真的必须遍历 resultFromA 并将其内容放入新的 Map 中吗？或者他们是我忽略的 API 中的方法/用法？

这是我想要做的：

OrtEnvironment environment = OrtEnvironment.getEnvironment();
OrtSession modelASession = environment.createSession(...);
OrtSession modelBSession = environment.createSession(...);

try (Result modelAResults = modelASession.run(inputTensor)) {
    
    try(Result modelBResults = modelBSession.run(modelAResults) { // <-------run won't take a Result object
        //do something with model B results...
    }
}

而我必须做的是：

OrtEnvironment environment = OrtEnvironment.getEnvironment();
OrtSession modelASession = environment.createSession(...);
OrtSession modelBSession = environment.createSession(...);

try (Result modelAResults = modelASession.run(inputTensor)) {

    Map<String,OnnxTensor> modelAMap = new HashMap<String,OnnxTensor>();
    modelAResults.forEach(e -> modelAMap.put(e.getKey(),(OnnxTensor) e.getValue()));
    
    try(Result modelBResults = modelBSession.run(modelAMap) {
        //do something with model B results...
    }
}

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）