Squeeze Onnx Import #1753

agelas · 2024-05-11T10:24:53Z

Pull Request Template

Checklist

Confirmed that run-checks all script has been executed.
Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#1714

Changes

Added the squeeze operation for to burn-import

Testing

onnx_tests
codegen_nodes test

codecov · 2024-05-11T21:52:19Z

Codecov Report

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

Project coverage is 86.41%. Comparing base (1073752) to head (3e60613).
Report is 7 commits behind head on main.

Files	Patch %	Lines
crates/burn-import/src/onnx/dim_inference.rs	65.51%	10 Missing ⚠️
crates/burn-import/src/onnx/op_configuration.rs	77.14%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1753      +/-   ##
==========================================
- Coverage   86.61%   86.41%   -0.21%     
==========================================
  Files         700      735      +35     
  Lines       83423    85729    +2306     
==========================================
+ Hits        72258    74083    +1825     
- Misses      11165    11646     +481

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

…ents

antimora

Thank you so much for completing a missing ONNX op in Burn. One down!

You've got almost perfect. We just need a couple fixes:

Add test case to verify axes by adding ONNX OpSet 13 (please see my inlined comments)
Account for negative values in axes.

antimora · 2024-05-13T16:47:04Z

crates/burn-import/src/onnx/op_configuration.rs

+        match key.as_str() {
+            "axes" => return value.clone().into_i64s(),
+            _ => {}


Great! This will support ONNX OPset 13

Can we, to be sure it works, add a unit test for OPset 13? This is similar to unsqueeze with opset 16 and 13.

Unsqueeze uses OPset 16 and 11. I think my original one used 16, so I added one for 13. Is that what you meant?

Yes, both should be good. OpSet 13 has axes attribute so it should work.

antimora · 2024-05-13T16:53:55Z

crates/burn-import/src/onnx/op_configuration.rs

+            }
+        }
+        _ => panic!("Arg for squeeze must be tensor or scalar"),
+    }


axes can contain negative values, which means counting dimensions from the back (see squeeze spec), and Burn squeeze only supports positive values (see doc). So we should account for this.

We are already doing this for one dimension (see code) for gather_config, so you can see the logic. We need to do this for all items in axes.

laggui

Thanks for your contribution! I have the same comments as already mentioned, and one question/comment regarding the squeeze op spec support.

This PR adds support for squeeze op on a single dim, so support is not complete but as long as it is explicit I have no issues adding full support if needed in another PR.

crates/burn-import/src/onnx/dim_inference.rs

laggui

Thanks for the changes! I have two minor comments and then it should be good to go! 🙂

crates/burn-import/onnx-tests/tests/squeeze/squeeze.onnx

laggui · 2024-05-14T12:06:06Z

crates/burn-import/src/onnx/dim_inference.rs

+    if let Some(Data::Int64s(axes)) = &node.inputs[1].value {
+        if axes.len() != 1 {
+            panic!("Squeeze: Only one axis should be specified for squeezing.");
+        }
+    } else {
+        panic!("Squeeze: Axes input must be an integer list.");
+    };


If we're checking for axes in the attributes in squeeze_config to support another opset we need to check here too to make sure the output is properly adjusted.

See the unsqueeze_update_output function which is doing something similar (though it captures the axes to support unsqueeze on multiple axes).

Sounds good, should be a little closer to what you had in mind now!

laggui

Great! Thanks for making the changes. I only have one minor comment left (we lost a check with the last committed change and I think we should keep it).

Should be good to go after that!

crates/burn-import/src/onnx/dim_inference.rs

laggui

Thanks for going through the changes!

LGTM 🚀

mtobin-tdab · 2024-07-17T13:17:21Z

Hi there, aware this is closed but wanted to ask why Squeeze has to specify only one axis? Is there a plan to make it possible to Squeeze to take > 1 axis? Should I make an issue to track this? In the mean time would it work to split into two layers which each Squeeze a different axis? Sorry I am not an expert in this field

  DEBUG onnx_ir::from_onnx: checking node squeeze1 for constants    
  DEBUG onnx_ir::from_onnx: checking input Argument { name: "const_axes__420", ty: Tensor(TensorType { elem_type: Int64, dim: 1, shape: Some([2]) }), value: Some(Int64s([2, 3])), passed: false } for const    
  ERROR burn_import::logger: PANIC => panicked at crates/onnx-ir/src/dim_inference.rs:463:9:
  Squeeze must specify only 1 axis, found 2

laggui · 2024-07-17T14:08:00Z

Hi @mtobin-tdab 👋

The ONNX spec is quite large so we don't necessarily aim to cover all the edge cases when adding support for an operator. That's why this PR was merged - we just make sure to panic when there's is a case which is not supported yet (such as yours).

You can definitely open an issue to ask for support for multiple axes. That's what issues are for (bugs, feature requests, etc.).

In the meantime, if you can generate your model with multiple squeeze nodes instead then yes that should work!

mtobin-tdab · 2024-07-17T14:15:17Z

Thank you @laggui !

agelas · 2024-07-17T18:46:32Z

@mtobin-tdab I think squeeze_dims added here in #1779 is what you might be looking for.

agelas · 2024-07-17T19:02:53Z

Ah wait, actually I think I see the issue. It looks like I didn't update the dimension inference, sorry about that!

laggui · 2024-07-17T19:16:29Z

Oh I didn't even see the follow-up PR for squeeze_dims 😄 looks like you've correctly identified the issue

agelas added 10 commits May 11, 2024 02:29

Fix typo

658add8

Add Python script for generating squeeze ONNX file

e1dbd3d

NodeCodegen for squeeze

dec8ba1

Add Squeeze to enum and macros

f5a046b

Add squeeze to LIFT_CONSTANT_FOR_NODE_TYPES

321ebb6

Squeeze configuration and node interpretation

160cac0

Include the ONNX model

d0ce463

Add squeeze test to onnx_tests

4763a8b

Add codegen test

c9547d3

Add squeeze_update_output

6b0b2d8

Update supported-onnx-ops file, cleanup comments, remove debug statem…

c05e0bf

…ents

agelas marked this pull request as ready for review May 13, 2024 02:08

nathanielsimard requested a review from laggui May 13, 2024 16:50

antimora requested changes May 13, 2024

View reviewed changes

laggui reviewed May 13, 2024

View reviewed changes

crates/burn-import/src/onnx/dim_inference.rs Outdated Show resolved Hide resolved

agelas added 4 commits May 13, 2024 11:44

Remove old squeeze.onnx

7643012

Test both opset 16 and 13

850367f

Make op config more spec compliant and add back squeeze.onnx

10f9fa4

Use clone_from to placate reviewdog

4410f72

laggui reviewed May 14, 2024

View reviewed changes

Better axes checking for squeeze in dim_inference

0da5788

laggui reviewed May 15, 2024

View reviewed changes

crates/burn-import/src/onnx/dim_inference.rs Show resolved Hide resolved

Check if single axes is specified

3e60613

laggui approved these changes May 15, 2024

View reviewed changes

agelas mentioned this pull request May 16, 2024

Contributor Book: Onnx to Burn Conversion #1771

Merged

2 tasks

nathanielsimard merged commit 9c5b07c into tracel-ai:main May 17, 2024
14 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Squeeze Onnx Import #1753

Squeeze Onnx Import #1753

agelas commented May 11, 2024 •

edited

Loading

codecov bot commented May 11, 2024 •

edited

Loading

antimora left a comment

antimora May 13, 2024

agelas May 13, 2024

antimora May 13, 2024

antimora May 13, 2024

laggui left a comment

laggui left a comment

laggui May 14, 2024

agelas May 14, 2024

laggui left a comment

laggui left a comment

mtobin-tdab commented Jul 17, 2024

laggui commented Jul 17, 2024

mtobin-tdab commented Jul 17, 2024

agelas commented Jul 17, 2024

agelas commented Jul 17, 2024

laggui commented Jul 17, 2024

Squeeze Onnx Import #1753

Squeeze Onnx Import #1753

Conversation

agelas commented May 11, 2024 • edited Loading

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

codecov bot commented May 11, 2024 • edited Loading

Codecov Report

antimora left a comment

Choose a reason for hiding this comment

antimora May 13, 2024

Choose a reason for hiding this comment

agelas May 13, 2024

Choose a reason for hiding this comment

antimora May 13, 2024

Choose a reason for hiding this comment

antimora May 13, 2024

Choose a reason for hiding this comment

laggui left a comment

Choose a reason for hiding this comment

laggui left a comment

Choose a reason for hiding this comment

laggui May 14, 2024

Choose a reason for hiding this comment

agelas May 14, 2024

Choose a reason for hiding this comment

laggui left a comment

Choose a reason for hiding this comment

laggui left a comment

Choose a reason for hiding this comment

mtobin-tdab commented Jul 17, 2024

laggui commented Jul 17, 2024

mtobin-tdab commented Jul 17, 2024

agelas commented Jul 17, 2024

agelas commented Jul 17, 2024

laggui commented Jul 17, 2024

agelas commented May 11, 2024 •

edited

Loading

codecov bot commented May 11, 2024 •

edited

Loading