[Feature Request] Python support for AddExternalInitializersFromFilesInMemory

[ardiloot]

Describe the feature request

There is AddExternalInitializersFromFilesInMemory method for loading external initializers from memory instead of a file on a disk in C/C++ api (docs).

However, this functionality seems to be not available from Python. When looking onnxruntime_pybind_state.cc, then there is add_external_initializers binding, but no add_external_initializers_from_files_in_memory.

Describe scenario use case

Adding add_external_initializers_from_files_in_memory allows loading models with external data directly from memory.

[xadupre]

Would you like to contribute?

[ardiloot]

Setting up the build tools on my side would involve quite a bit of overhead, so it’s not my preferred approach. If needed, I can prepare a PR as a last resort, but I’d prefer to explore other options first

[xadupre]

If you have a unit test or an example in mind just to make sure I implement what you need, that would be great. Then I'll do the rest.

[ardiloot]

Thank you! I composed an example:

import onnx
import numpy as np
import onnxruntime as ort
from pathlib import Path
from onnx.external_data_helper import convert_model_to_external_data

if __name__ == "__main__":
    test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
    original_model_path = "mobilenetv3_small_075_Opset17.onnx"
    # https://github.com/onnx/models/blob/main/Computer_Vision/mobilenetv3_small_075_Opset17_timm/mobilenetv3_small_075_Opset17.onnx

    # Convert model to use external data for weights
    # Results in model_external_data.onnx (52KB) and weights.bin (7.8MB)
    onnx_model = onnx.load(original_model_path)
    convert_model_to_external_data(
        onnx_model,
        all_tensors_to_one_file=True,
        location="weights.bin",
    )
    onnx.save_model(onnx_model, "model_external_data.onnx")

    # Run inference with the original model
    options1 = ort.SessionOptions()
    sess1 = ort.InferenceSession(original_model_path, options=options1)
    outputs1 = sess1.run(None, {"x": test_input})[0]
    print(f"Output #1 shape: {outputs1.shape}, mean: {outputs1.mean()}")

    # Run inference with the external data model (from disk)
    options2 = ort.SessionOptions()
    sess2 = ort.InferenceSession("model_external_data.onnx", options=options2)
    outputs2 = sess2.run(None, {"x": test_input})[0]
    print(f"Output #2 shape: {outputs2.shape}, mean: {outputs2.mean()}")
    np.testing.assert_equal(outputs1, outputs2)

    # Run inference with the external data model (from memory)
    model_data = open("model_external_data.onnx", "rb").read()
    weights_data = open("weights.bin", "rb").read()
    # Remove the weights file after reading to make sure it's not used
    Path("weights.bin").unlink(missing_ok=True)
    options3 = ort.SessionOptions()
    # This is currently not working due to missing add_external_initializers_from_files_in_memory()
    options3.add_external_initializers_from_files_in_memory(
        {"weights.bin": weights_data}
    )
    sess3 = ort.InferenceSession(model_data, options=options3)
    outputs3 = sess3.run(None, {"x": test_input})[0]
    print(f"Output #3 shape: {outputs3.shape}, mean: {outputs3.mean()}")
    np.testing.assert_equal(outputs1, outputs3)