[code-reflection] RFR: Model lifetimes of onnx session-related objects more explicitly [v4]
Maurizio Cimadamore
mcimadamore at openjdk.org
Mon Mar 3 22:01:53 UTC 2025
On Mon, 3 Mar 2025 21:58:57 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.
>
> Maurizio Cimadamore has refreshed the contents of this pull request, and previous commits have been removed. The incremental views will show differences compared to the previous content of the PR. The pull request contains one new commit since the last revision:
>
> Merge branch 'code-reflection' into onnx_session_lifetime
cr-examples/onnx/src/test/java/oracle/code/onnx/mnist/MNISTDemo.java line 111:
> 109: public static float[] classify(float[] imageData) {
> 110: try (Arena arena = Arena.ofConfined()) {
> 111: var imageTensor = Tensor.ofShape(new long[]{1, 1, IMAGE_SIZE, IMAGE_SIZE}, imageData);
I believe it would be useful to have overloads of the tensor factories taking an arena...
-------------
PR Review Comment: https://git.openjdk.org/babylon/pull/332#discussion_r1978303702
More information about the babylon-dev
mailing list