[code-reflection] RFR: Model lifetimes of onnx session-related objects more explicitly [v4]
Maurizio Cimadamore
mcimadamore at openjdk.org
Mon Mar 3 22:01:52 UTC 2025
> 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
-------------
Changes:
- all: https://git.openjdk.org/babylon/pull/332/files
- new: https://git.openjdk.org/babylon/pull/332/files/0605aa4b..b7fa2555
Webrevs:
- full: https://webrevs.openjdk.org/?repo=babylon&pr=332&range=03
- incr: Webrev is not available because diff is too large
Stats: 341809 lines in 3274 files changed: 210309 ins; 103791 del; 27709 mod
Patch: https://git.openjdk.org/babylon/pull/332.diff
Fetch: git fetch https://git.openjdk.org/babylon.git pull/332/head:pull/332
PR: https://git.openjdk.org/babylon/pull/332
More information about the babylon-dev
mailing list