Skip to content

Commit

Permalink
Separating ONNX parsing from burn-import (#1921)
Browse files Browse the repository at this point in the history
* separating onnx parsing from burn-import

* ran clippy and cargo-fmt

* removed unused deps from onnx-ir

* fixed clippy warnings that were causing run-checks to fail

* removed dead code

* removed unused dependencies from burn-import

* updated contributor-book, updated publish.yml, added readme

* update cargo lock

* formatted md document with prettier, rephrased sentence

* missed the errors with reduce_prod_conversion during merge

* formatted onnx-to-burn-conversion-tool.md, forgot to save
  • Loading branch information
skewballfox authored Jul 2, 2024
1 parent 755c070 commit 25348cf
Show file tree
Hide file tree
Showing 22 changed files with 405 additions and 279 deletions.
6 changes: 6 additions & 0 deletions .github/workflows/publish.yml
Original file line number Diff line number Diff line change
Expand Up @@ -161,3 +161,9 @@ jobs:
with:
crate: burn-import
secrets: inherit

publish-onnx-ir:
uses: tracel-ai/burn/.github/workflows/publish-template.yml@main
with:
crate: onnx-ir
secrets: inherit
23 changes: 18 additions & 5 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

28 changes: 18 additions & 10 deletions contributor-book/src/guides/adding-a-new-operation-to-burn.md
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,8 @@ plug in your operator in terms of \\(x\\) and \\(y\\), and just swap out the var

### Testing autodiff

For testing the `autodiff` operations, please refer to [this section](../getting-started/testing.md).
For testing the `autodiff` operations, please refer to
[this section](../getting-started/testing.md).

## Adding the Op to other backends

Expand Down Expand Up @@ -199,11 +200,15 @@ Generating the ONNX test files or tests is already covered
about the specific changes you need to make when adding new operators after you have generated the
tests.

The crate is divided into two sections `src/burn` and `src/onnx`. The code under the former
corresponds to the operation you've implemented earlier in this guide, and the latter to the
operations defined in the ONNX specification. So when you are loading a model, the operator is first
parsed to an intermediate representation defined by `src/onnx`, and then mapped to a Burn operation
defined under `src/burn/node`.
Changes will need to be made to both `onnx-ir` and `burn-import`. The code within `onnx-ir` defines
how to parse the nodes in an onnx file and produces the intermediate representation. The code within
`burn-import` is divided into two sections: `src/onnx` and `src/burn`. The code under the former
maps that intermediate representation to one used for code generation and the latter defines how to
generate code for the operator you've implemented earlier in this guide.

So when you are loading a model, the operator is first parsed to an intermediate representation
defined by `burn-import` and then mapped to a Burn operation defined under `src/burn/node`; the
mapping from onnx to burn is aptly defined in `src/onnx/to_burn`

Let's review the changes made for powf starting from `src/burn` and moving to `src/onnx`:

Expand All @@ -218,17 +223,20 @@ Let's review the changes made for powf starting from `src/burn` and moving to `s
[`{op}_conversion` function](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-import/src/onnx/to_burn.rs#L717)
that maps the ONNX node to the binary type
3. Specify how dimensions for the output should be derived in
[crates/burn-import/src/onnx/dim_inference.rs](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-import/src/onnx/dim_inference.rs#L55)
[crates/onnx-ir/src/dim_inference.rs](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/dim_inference.rs#L17)

And you're done! Congrats, you just fully added a new operation to burn, and we are all one step
closer to the answer to [Are we learning yet?](https://www.arewelearningyet.com/) being "Yes, and
it's freaking fast!". Buy yourself a coffee.

[^supertrait]: for more on supertraits see
[^supertrait]:
for more on supertraits see
[the advanced trait section of the rust book](https://doc.rust-lang.org/book/ch19-03-advanced-traits.html#using-supertraits-to-require-one-traits-functionality-within-another-trait)

[^autodiff]: wiki link for
[^autodiff]:
wiki link for
[automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation)

[^absolute_units]: for more information on unit structs see
[^absolute_units]:
for more information on unit structs see
[the defining and instantiating structs section of the rust book](https://doc.rust-lang.org/book/ch05-01-defining-structs.html#unit-like-structs-without-any-fields)
28 changes: 27 additions & 1 deletion contributor-book/src/guides/onnx-to-burn-conversion-tool.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,17 @@ For an introduction to ONNX import in Burn, see
- [Design Goals](#design-goals)
- [Design Decisions](#design-decisions)
- [Adding New Operators](#adding-new-operators)
- [Implementing a New Operator](#implementing-a-new-operator)
- [Implementing a New Operator](#implementing-a-new-operator)
- [Step 1: Visibility](#step-1-visibility)
- [Step 2: Node Implementation](#step-2-node-implementation)
- [Within Onnx-IR](#within-onnx-ir)
- [Within burn-import](#within-burn-import)
- [Step 3: Registering New Operations](#step-3-registering-new-operations)
- [Step 4: Create a Config Function](#step-4-create-a-config-function)
- [Step 5: Dimension Inference](#step-5-dimension-inference)
- [Step 6: Integrate into the Graph Building Process](#step-6-integrate-into-the-graph-building-process)
- [Step 7: Add Newly Supported Op!](#step-7-add-newly-supported-op)
- [Misc:](#misc)
- [Testing](#testing)
- [Resources](#resources)

Expand Down Expand Up @@ -91,6 +101,22 @@ located in the `src/burn/node/` directory.

### Step 2: Node Implementation

#### Within Onnx-IR

If the node type does not exist within the
[`NodeType` enum](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/ir.rs#L246),
it will need to be added (support for custom operators is planned). If the node might be provided an
input which is a constant or the output of an identity node, it will need to be added to the list of
nodeTypes
[checked for constants](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/from_onnx.rs#L21).
The node will need to be added to `dim_inference`, and in most cases the work parsing side will be
done. If a node requires extra parsing (such as handling an edge case like potentially remapping an
unsqueeze to a reshape) the best place for that is after check constants and prior to dim_inference
in
[`OnnxGraphBuilder::Build`](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/from_onnx.rs#L221)

#### Within burn-import

Create a new file named `<operation_name>.rs` in the `src/burn/node/` directory.
This file will define the structure and functionality of your new operation. By convention, the
necessary information for carrying out an operation is encapsulated within a struct named
Expand Down
9 changes: 1 addition & 8 deletions crates/burn-import/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,30 +20,23 @@ pytorch = ["burn/record-item-custom-serde", "thiserror", "zip"]

[dependencies]
burn = { path = "../burn", version = "0.14.0", features = ["ndarray"] }

bytemuck = { workspace = true }
onnx-ir = { path = "../onnx-ir" }
candle-core = { workspace = true }
derive-new = { workspace = true }
half = { workspace = true }
log = { workspace = true }
proc-macro2 = { workspace = true }
protobuf = { workspace = true, features = ["with-bytes"] }
quote = { workspace = true }
regex = { workspace = true }
rust-format = { workspace = true, features = ["token_stream", "post_process"] }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true, features = ["std"] }
strum = { workspace = true }
strum_macros = { workspace = true }
syn = { workspace = true, features = ["parsing"] }
thiserror = { workspace = true, optional = true }
tracing-core = { workspace = true }
tracing-subscriber = { workspace = true }
zip = { workspace = true, optional = true }

[build-dependencies]
protobuf-codegen = { workspace = true }

[dev-dependencies]
pretty_assertions = { workspace = true }
rstest = { workspace = true }
11 changes: 0 additions & 11 deletions crates/burn-import/build.rs

This file was deleted.

11 changes: 0 additions & 11 deletions crates/burn-import/src/onnx/mod.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,3 @@
mod coalesce;
mod dim_inference;
mod from_onnx;
mod ir;
mod node_remap;
mod op_configuration;
mod proto_conversion;
mod protos;
mod to_burn;

pub use to_burn::*;

pub use from_onnx::parse_onnx;
pub use ir::OnnxGraph;
2 changes: 1 addition & 1 deletion crates/burn-import/src/onnx/op_configuration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ use burn::nn::{
PaddingConfig2d,
};

use super::ir::{ArgType, AttributeValue, Data, Node};
use crate::burn::node::resize::ResizeMode;
use onnx_ir::ir::{ArgType, AttributeValue, Data, Node};

/// Create a Conv1dConfig from the attributes of the node
pub fn conv1d_config(curr: &Node) -> Conv1dConfig {
Expand Down
Loading

0 comments on commit 25348cf

Please sign in to comment.