Add support for OnnxRuntime Profiler (#2472)

deepjavalibrary · Mar 23, 2023 · fb304a3 · fb304a3
1 parent 9655c12
commit fb304a3
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/docs/development/profiler.md b/docs/development/profiler.md
@@ -60,3 +60,20 @@ Each record contains `name`(operator name), `dur`(time duration), `shape`(input
   "args": {}
 }
 ```
+
+### OnnxRuntime
+
+When loading a model, the profiler can be enabled by specifying the desired filepath in the criteria:
+
+```
+    Criteria<Image, Classifications> criteria =
+        Criteria.builder()
+            .optOption("profilerOutput", "build/testOrtProfiling")
+            ...
+            .build()
+```
+
+Following this, the file will appear at "build/tesetOrtProfiling_YYYY-MM-DD.json".
+
+You can view it in a browser using a trace consumer like `chrome://tracing `.
+For more details, see the [OnnxRuntime Profiler Docs](https://onnxruntime.ai/docs/performance/tune-performance/profiling-tools.html).
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
@@ -192,6 +192,11 @@ private SessionOptions getSessionOptions(Map<String, ?> options) throws OrtExcep
             ortSession.registerCustomOpLibrary(customOpLibrary);
         }
 
+        String profilerOutput = (String) options.get("profilerOutput");
+        if (profilerOutput != null) {
+            ortSession.enableProfiling(profilerOutput);
+        }
+
         Device device = manager.getDevice();
         if (options.containsKey("ortDevice")) {
             String ortDevice = (String) options.get("ortDevice");

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
@@ -60,6 +60,7 @@ public void testOrt() throws TranslateException, ModelException, IOException {
                             .optOption("memoryPatternOptimization", "true")
                             .optOption("cpuArenaAllocator", "true")
                             .optOption("disablePerSessionThreads", "true")
+                            .optOption("profilerOutput", "build/testOrtProfiling")
                             .build();
 
             IrisFlower virginica = new IrisFlower(1.0f, 2.0f, 3.0f, 4.0f);