Support set schema inference function in python #5940

OYCN · 2024-02-16T07:40:38Z

Description

Support set schema inference callback function for custion OpSchema before register in python.

example:

def func(ctx: onnx.shape_inference.InferenceContext):
    # get some info
    assert ctx.get_num_inputs() == 2
    value = ctx.get_input_type(0)
    ...
    # get or create output proto object
    output = ctx.get_output_type(0)
    # set type or shape
    ...
    # set the result proto
    ctx.set_output_type(0, output)

schema.set_type_and_shape_inference_function(func)

Note

Depends on #5906

Motivation and Context

Follow up of #5019

Signed-off-by: opluss <[email protected]>

OYCN · 2024-02-16T07:45:42Z

We're spending a considerable amount of code on passing proto objects between C++ and Python. However, the repository seems well-prepared for this feature.

Signed-off-by: opluss <[email protected]>

codecov · 2024-02-16T07:57:51Z

Codecov Report

Attention: Patch coverage is 93.54839% with 4 lines in your changes are missing coverage. Please review.

Project coverage is 57.10%. Comparing base (83194ed) to head (73bed3b).
Report is 35 commits behind head on main.

Files	Patch %	Lines
onnx/shape_inference.py	83.33%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5940      +/-   ##
==========================================
+ Coverage   56.95%   57.10%   +0.14%     
==========================================
  Files         506      506              
  Lines       30467    30989     +522     
  Branches     4592     4601       +9     
==========================================
+ Hits        17353    17696     +343     
- Misses      12285    12467     +182     
+ Partials      829      826       -3

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

onnx/cpp2py_export.cc

Signed-off-by: oPluss <[email protected]>

add check in `InferenceContext` py wrapper Signed-off-by: opluss <[email protected]>

Signed-off-by: oPluss <[email protected]>

Signed-off-by: opluss <[email protected]>

OYCN · 2024-03-29T10:34:49Z

Hi @justinchuby ,

In this PR, we can implement shape inference on the Python side, similar to how it's done on the C++ side. If you have any suggestions for this implementation, I'm open to making adjustments accordingly.

justinchuby · 2024-03-29T18:38:59Z

Thank you! Is it ready to be reviewed?

OYCN · 2024-03-31T06:43:27Z

Thank you! Is it ready to be reviewed?

Yes, I have removed the 'WIP' prefix from the title. Please feel free to leave any comments. :D

xadupre · 2024-04-09T07:43:07Z

onnx/test/shape_inference_test.py

+ )
+
+ assert ctx.get_num_inputs() == 2
+ in0 = ctx.get_input_type(0)


My concern with this is it goes through serialization to access the type information. It is not really efficient. I would change the API so that it does not return a TypeProto but the type and the shape as regular python objects.

I think I agree with Xavier, but I am a bit confused also. I see the method implementation serializes proto values to string and returns them. We could just return a pointer to the C++ Proto object (wrapped as a Python object). Is that your suggestion Xavier?

Should we closely mimic the C++ API design interface, or should we integrate Python's native types for interactions? Utilizing the proto pointer for interactions may require additional codes to bind them to Python (If there is another way please correct me), or we need to include some third-party library.

xadupre · 2024-04-09T07:44:43Z

onnx/test/shape_inference_test.py

+ out.tensor_type.shape.dim.add().dim_value = N
+ out.tensor_type.shape.dim.add().dim_value = La * Lb
+ out.tensor_type.shape.dim.add().dim_value = out_len[i]
+ ctx.set_output_type(i, out)


Same comment here, we should avoid serialization with something like set_output_type_and_shape(in0_type, (N, La*Lb, out_lin[i]). The type is created on C++ side, there is no serialization and it would be more efficient.

xadupre · 2024-04-09T07:45:42Z

It would be a nice feature to have.

gramalingam · 2024-04-17T20:02:37Z

onnx/cpp2py_export.cc

@@ -114,6 +116,9 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
 #endif // ONNX_ML
 );

+ // Avoid Segmentation fault if we not free the python function in Custom Schema
+ onnx_cpp2py_export.add_object("_cleanup", py::capsule([] { OpSchemaRegistry::OpSchemaDeregisterAll(); }));


Can you clarify when this gets invoked?

I think the segfault is caused by:
The python object (the inference function in custom schema) need destroyed before the python interpreter is destroyed. The static container within the schema factory is destroyed after main function and before the interpreter. Therefore, we need to manually destroy the Python object.

About '_cleanup' : https://pybind11.readthedocs.io/en/stable/advanced/misc.html#module-destructors

onnx/cpp2py_export.cc

gramalingam · 2024-04-17T20:44:49Z

Thanks for creating the PR! It would be great to add this functionality. My comments above.

OYCN

Forgive my mistake, I didn't know that the response need to be commit manually. My reply is always in pending status.

OYCN · 2024-03-29T10:31:41Z

onnx/cpp2py_export.cc

@@ -114,6 +116,9 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
 #endif // ONNX_ML
 );

+ // Avoid Segmentation fault if we not free the python function in Custom Schema
+ onnx_cpp2py_export.add_object("_cleanup", py::capsule([] { OpSchemaRegistry::OpSchemaDeregisterAll(); }));


Perhaps a more graceful approach would be to collect the schemas registered from Python and deregister them during cleanup. However, I'm not sure if it's worth the effort. In most cases, invoking cleanup implies that Python is exiting.

OYCN · 2024-04-28T00:28:14Z

onnx/test/shape_inference_test.py

+ )
+
+ assert ctx.get_num_inputs() == 2
+ in0 = ctx.get_input_type(0)


Should we closely mimic the C++ API design interface, or should we integrate Python's native types for interactions? Utilizing the proto pointer for interactions may require additional codes to bind them to Python (If there is another way please correct me), or we need to include some third-party library.

OYCN · 2024-05-24T14:50:58Z

onnx/cpp2py_export.cc

@@ -114,6 +116,9 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
 #endif // ONNX_ML
 );

+ // Avoid Segmentation fault if we not free the python function in Custom Schema
+ onnx_cpp2py_export.add_object("_cleanup", py::capsule([] { OpSchemaRegistry::OpSchemaDeregisterAll(); }));


I think the segfault is caused by:
The python object (the inference function in custom schema) need destroyed before the python interpreter is destroyed. The static container within the schema factory is destroyed after main function and before the interpreter. Therefore, we need to manually destroy the Python object.

About '_cleanup' : https://pybind11.readthedocs.io/en/stable/advanced/misc.html#module-destructors

xadupre · 2024-05-27T08:07:43Z

I would do this in a different way. I think we should export in python all functions doing shape inference for every operator. We should have a python version of the algorithm doing shape inference. That way, adding new shape inference function would be easier and we could do local shape inference. This feature is needed when building a model. Shapes can be used to write efficient sequences of onnx operators.

OYCN · 2024-05-30T01:21:38Z

Hi Xavier,

After thinking about it for a while, please correct me if my understanding is wrong. Implementing a Python version of shape inference may not eliminate the need to pass ONNX protos between C++ and Python. In the current architecture, if we want to obtain any proto in the shape inference function, we need to export InferenceContext from C++ to Python. Perhaps a more efficient implementation for interacting with information carried by proto is key to solving the problem. This could be done using native Python types (though it’s challenging to deliver graph protos), type pointers, or other methods.

Best Regards.

Signed-off-by: opluss <[email protected]>

add schema inference py interface

d6d944a

Signed-off-by: opluss <[email protected]>

fix type error

e813a3a

Signed-off-by: opluss <[email protected]>

Merge branch 'onnx:main' into dev-py-schema-inference

b558aed

github-advanced-security bot found potential problems Feb 26, 2024

View reviewed changes

onnx/cpp2py_export.cc Fixed Show resolved Hide resolved

onnx/cpp2py_export.cc Fixed Show resolved Hide resolved

onnx/cpp2py_export.cc Fixed Show resolved Hide resolved

OYCN added 7 commits February 28, 2024 14:06

fix bug

095f0fd

Signed-off-by: oPluss <[email protected]>

Merge branch 'main' into dev-py-schema-inference

a3a00d6

add unittest for custom schema shape inference

b47c69e

add check in `InferenceContext` py wrapper Signed-off-by: opluss <[email protected]>

Merge branch 'main' into dev-py-schema-inference

61a5305

Signed-off-by: oPluss <[email protected]>

add check for InferenceContext py wrapper

f9beedb

Signed-off-by: opluss <[email protected]>

remove logic because it was in another pr

3132b4d

Signed-off-by: opluss <[email protected]>

Merge branch 'main' into dev-py-schema-inference

3a57147

justinchuby marked this pull request as ready for review March 29, 2024 18:38

justinchuby requested a review from a team as a code owner March 29, 2024 18:38

OYCN changed the title ~~[WIP] Support set schema inference function in python~~ Support set schema inference function in python Mar 31, 2024

justinchuby self-assigned this Apr 1, 2024

Merge branch 'main' into dev-py-schema-inference

e3b744b

xadupre reviewed Apr 9, 2024

View reviewed changes

Merge branch 'main' into dev-py-schema-inference

01f955d

gramalingam reviewed Apr 17, 2024

View reviewed changes

onnx/cpp2py_export.cc Outdated Show resolved Hide resolved

gramalingam reviewed Apr 17, 2024

View reviewed changes

onnx/cpp2py_export.cc Outdated Show resolved Hide resolved

gramalingam reviewed Apr 17, 2024

View reviewed changes

onnx/cpp2py_export.cc Outdated Show resolved Hide resolved

OYCN commented May 24, 2024

View reviewed changes

akmaru mentioned this pull request May 25, 2024

[CustomOp] Support to registor OpSchema of OrtCustomOp akmaru/onnxruntime#1

Open

return None if target is not exist

73bed3b

Signed-off-by: opluss <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support set schema inference function in python #5940

Support set schema inference function in python #5940

OYCN commented Feb 16, 2024 •

edited

OYCN commented Feb 16, 2024

codecov bot commented Feb 16, 2024 •

edited

OYCN commented Mar 29, 2024

justinchuby commented Mar 29, 2024

OYCN commented Mar 31, 2024

xadupre Apr 9, 2024

gramalingam Apr 17, 2024

OYCN Apr 28, 2024

xadupre Apr 9, 2024

xadupre commented Apr 9, 2024 •

edited by justinchuby

gramalingam Apr 17, 2024

OYCN May 24, 2024

gramalingam commented Apr 17, 2024

OYCN left a comment

OYCN Mar 29, 2024

OYCN Apr 28, 2024

OYCN May 24, 2024

xadupre commented May 27, 2024

OYCN commented May 30, 2024

Support set schema inference function in python #5940

Are you sure you want to change the base?

Support set schema inference function in python #5940

Conversation

OYCN commented Feb 16, 2024 • edited

Description

Note

Motivation and Context

OYCN commented Feb 16, 2024

codecov bot commented Feb 16, 2024 • edited

Codecov Report

OYCN commented Mar 29, 2024

justinchuby commented Mar 29, 2024

OYCN commented Mar 31, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xadupre commented Apr 9, 2024 • edited by justinchuby

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gramalingam commented Apr 17, 2024

OYCN left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xadupre commented May 27, 2024

OYCN commented May 30, 2024

OYCN commented Feb 16, 2024 •

edited

codecov bot commented Feb 16, 2024 •

edited

xadupre commented Apr 9, 2024 •

edited by justinchuby