onnxruntime icon indicating copy to clipboard operation
onnxruntime copied to clipboard

[Mobile] Bug in android implementation for loading large onnx models

Open l3utterfly opened this issue 1 year ago • 2 comments

Describe the issue

The current code tries to load the file on the JVM heap, before passing it to the native code, resulting in unable to load large onnx files on Android due to limited jvm heap.

Can be fixed by commenting out the load from assets section in OnnxruntimeModule.java:

// load model via model path string uri
      // InputStream modelStream =
      //     reactContext.getApplicationContext().getContentResolver().openInputStream(Uri.parse(uri));
      // Reader reader = new BufferedReader(new InputStreamReader(modelStream));
      // byte[] modelArray = new byte[modelStream.available()];
      // modelStream.read(modelArray);
      // modelStream.close();

      // directly pass the uri to createSession, we don't support loading from assets
      ortSession = ortEnvironment.createSession(uri, sessionOptions);

This disables loading from assets. Perhaps consider adding more logic here that bypasses the JVM heap if the uri is from disk.

To reproduce

Load any large onnx file greater than 500MB on Android

Urgency

No response

Platform

Android

OS Version

Android 14

ONNX Runtime Installation

Built from Source

Compiler Version (if 'Built from Source')

No response

Package Name (if 'Released Package')

None

ONNX Runtime Version or Commit ID

1.16

ONNX Runtime API

Java

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

l3utterfly avatar Feb 22 '24 06:02 l3utterfly

@Craigacp any thoughts?

Feels like duplicating the model bytes in memory is coming up a few times lately with models getting much bigger. I'm definitely not a Java expert so I don't know what the best approach is to avoid that happening.

If ORT can read from a path on Android with no additional logic that might be the best general solution to recommend/implement.

IIRC passing raw bytes in is ineffective in reducing memory as we have to call https://github.com/microsoft/onnxruntime/blob/ae3d73c9818c34af42c785ff2bd9558007ba315f/onnxruntime/core/session/inference_session.cc#L630 anyway as the bytes are protobuf encoded.

skottmckay avatar Feb 23 '24 09:02 skottmckay

I can add a Java session constructor which accepts a ByteBuffer, then users could memory map the file on disk with (https://docs.oracle.com/javase/8/docs/api/java/nio/channels/FileChannel.html#map-java.nio.channels.FileChannel.MapMode-long-long-), pass through the ByteBuffer to JNI where I can pull out the address from it and then hand that byte array into the C API session constructor. However I've not done much work with memory mapping in Java, so I don't know if there are additional considerations or if it's troublesome on Android. I also don't know if that would help the react native wrapper as I don't know anything about react native.

I agree that if the native code can read from a filesystem path on Android then we should expose that too, but it's already exposed in Java & Android so maybe that's just a react native problem?

Craigacp avatar Feb 23 '24 15:02 Craigacp

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

github-actions[bot] avatar Mar 25 '24 15:03 github-actions[bot]