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 <opluss@qq.com>

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 <opluss@qq.com>

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.11%. Comparing base (83194ed) to head (01f955d).
Report is 13 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.11%   +0.15%     
==========================================
  Files         506      506              
  Lines       30467    30544      +77     
  Branches     4592     4602      +10     
==========================================
+ Hits        17353    17445      +92     
+ Misses      12285    12273      -12     
+ 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 <opluss@qq.com>

add check in `InferenceContext` py wrapper Signed-off-by: opluss <opluss@qq.com>

Signed-off-by: oPluss <opluss@qq.com>

Signed-off-by: opluss <opluss@qq.com>

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?

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?

gramalingam · 2024-04-17T20:15:19Z

onnx/cpp2py_export.cc

+    }
+    auto attr = ctx->getAttribute(name);
+    if (attr == nullptr) {
+      fail_shape_inference("Internal error: `attr` is nullptr in `__get_attribute`");


This is not an error situation. In C++, this single method plays a dual role, both to check if an attribute is present, and to get its value. Is it possible to return None in this situation?

Or, just a return a C++ pointer (using pybind's conventions ... it must have some standard approach).

gramalingam · 2024-04-17T20:27:01Z

onnx/cpp2py_export.cc

+    }
+    auto tensor = ctx->getInputData(index);
+    if (tensor == nullptr) {
+      fail_shape_inference("Internal error: `tensor` is nullptr in `__get_input_data`");


Again, not an error. Signifies no static-constant-value is known for this input. Best to map this to None or some similar thing

gramalingam · 2024-04-17T20:27:21Z

onnx/cpp2py_export.cc

+    }
+    auto stensor = ctx->getInputSparseData(index);
+    if (stensor == nullptr) {
+      fail_shape_inference("Internal error: `stensor` is nullptr in `__get_input_sparse_data`");


As above. Map to None?

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

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

add schema inference py interface

d6d944a

Signed-off-by: opluss <opluss@qq.com>

fix type error

e813a3a

Signed-off-by: opluss <opluss@qq.com>

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 <opluss@qq.com>

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 <opluss@qq.com>

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

61a5305

Signed-off-by: oPluss <opluss@qq.com>

add check for InferenceContext py wrapper

f9beedb

Signed-off-by: opluss <opluss@qq.com>

remove logic because it was in another pr

3132b4d

Signed-off-by: opluss <opluss@qq.com>

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

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

xadupre Apr 9, 2024

xadupre commented Apr 9, 2024 •

edited by justinchuby

gramalingam Apr 17, 2024

gramalingam Apr 17, 2024

gramalingam Apr 17, 2024

gramalingam Apr 17, 2024

gramalingam Apr 17, 2024

gramalingam commented Apr 17, 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

xadupre commented Apr 9, 2024 • edited by justinchuby

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

Choose a reason for hiding this comment

gramalingam commented Apr 17, 2024

OYCN commented Feb 16, 2024 •

edited

codecov bot commented Feb 16, 2024 •

edited

xadupre commented Apr 9, 2024 •

edited by justinchuby