From ac4ec939614bdb65ef74c0842d145e746c0d01fc Mon Sep 17 00:00:00 2001 From: andreaduque Date: Wed, 23 Feb 2022 09:50:46 +0100 Subject: [PATCH 1/3] arena allocator setter --- .../src/main/java/ai/djl/onnxruntime/engine/OrtModel.java | 6 ++++++ .../src/test/java/ai/djl/onnxruntime/engine/OrtTest.java | 1 + 2 files changed, 7 insertions(+) diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java index c3d856a83bc..66dda422410 100644 --- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java +++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java @@ -180,6 +180,12 @@ private SessionOptions getSessionOptions(Map options) throws OrtExcep if (Boolean.parseBoolean(memoryOptimization)) { ortSession.setMemoryPatternOptimization(true); } + + String cpuArena = (String) options.get("memoryPatternOptimization"); + if (Boolean.parseBoolean(cpuArena)) { + ortSession.setCPUArenaAllocator(true); + } + return ortSession; } } diff --git a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java index 94fab66b015..6d6c709f03c 100644 --- a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java +++ b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java @@ -48,6 +48,7 @@ public void testOrt() throws TranslateException, ModelException, IOException { .optOption("executionMode", "SEQUENTIAL") .optOption("optLevel", "NO_OPT") .optOption("memoryPatternOptimization", "true") + .optOption("cpuArenaAllocator", "true") .build(); IrisFlower virginica = new IrisFlower(1.0f, 2.0f, 3.0f, 4.0f); From 4c750e10d017fbeb1a81ae711c3423e6c6616d13 Mon Sep 17 00:00:00 2001 From: andreaduque Date: Thu, 24 Feb 2022 09:28:20 +0100 Subject: [PATCH 2/3] fix options name --- .../src/main/java/ai/djl/onnxruntime/engine/OrtModel.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java index 66dda422410..4cf815fe98c 100644 --- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java +++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java @@ -181,7 +181,7 @@ private SessionOptions getSessionOptions(Map options) throws OrtExcep ortSession.setMemoryPatternOptimization(true); } - String cpuArena = (String) options.get("memoryPatternOptimization"); + String cpuArena = (String) options.get("cpuArenaAllocator"); if (Boolean.parseBoolean(cpuArena)) { ortSession.setCPUArenaAllocator(true); } From a5e48a8bd579ab43db28095e6ae8c4c5502179e7 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Thu, 24 Feb 2022 09:22:24 -0800 Subject: [PATCH 3/3] [onnx] Fixes java code format Change-Id: I9c85e84f420ae02ca86f6f9856e00b20186fbfb6 --- .../src/main/java/ai/djl/onnxruntime/engine/OrtModel.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java index 4cf815fe98c..dfd9caa1fbc 100644 --- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java +++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java @@ -184,7 +184,7 @@ private SessionOptions getSessionOptions(Map options) throws OrtExcep String cpuArena = (String) options.get("cpuArenaAllocator"); if (Boolean.parseBoolean(cpuArena)) { ortSession.setCPUArenaAllocator(true); - } + } return ortSession; }