0

我有两个使用 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...
    }
}
4

0 回答 0