[code-reflection] RFR: Model lifetimes of onnx session-related objects more explicitly
Adam Sotona
asotona at openjdk.org
Sun Mar 2 15:25:09 UTC 2025
On Fri, 28 Feb 2025 12:42:24 GMT, Maurizio Cimadamore <mcimadamore at openjdk.org> wrote:
> The class representing an onnx session is auto closeable. But, in the current code, a session is closed immediately after its `run` method is called. This is problematic because a session returns some ORTValues (tensors) which also need to be freed, but that cannot be freed immediately after calling `run` (as they need to be used by clients).
>
> To address this problem, I tweaked the session code to accept an external arena. All the allocation of session-related data structures now happens using that external arena. This means that the client can now be in charge of managing the lifetime of a session (see changes to MNIST demo).
>
> To test, I tweaked the MNIST code to do 10K iterations on each button pressed. Predictably, a single button pressed resulted in over 3g of memory being leaked. With these changes the memory arrives at ~400K (there is still some minor leak, but not sure worth pushing more).
>
> If the changes to the demo are not deemed good, I can withdraw this PR -- I mostly wanted to capture the result of my exploration somewhere.
cr-examples/onnx/src/test/java/oracle/code/onnx/MNISTDemo.java line 151:
> 149: var imageTensor = Tensor.ofShape(new long[]{1, 1, IMAGE_SIZE, IMAGE_SIZE}, imageData);
> 150:
> 151: try (Arena onnxSession = Arena.ofConfined()) {
Please don't call it a session. Onnx session is created below during the execution.
cr-examples/onnx/src/test/java/oracle/code/onnx/RuntimeTest.java line 18:
> 16: public void test() throws Exception {
> 17: var ort = OnnxRuntime.getInstance();
> 18: try (Arena sessionArena = Arena.ofConfined()) {
The same naming confusion is here. Below are created two sessions.
-------------
PR Review Comment: https://git.openjdk.org/babylon/pull/332#discussion_r1976652470
PR Review Comment: https://git.openjdk.org/babylon/pull/332#discussion_r1976652723
More information about the babylon-dev
mailing list