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

Question

我有两个使用 ONNX 构建的模型，模型 A 和模型 B。我使用ONNX 运行时 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...
    }
}

标签： javaonnxonnxruntime