Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ONNX debug improvements #1712

Merged
merged 3 commits into from
Apr 30, 2024
Merged

ONNX debug improvements #1712

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b3513a0
Minor debug improvements
antimora Apr 30, 2024
2b7e45e
Change warn to panic
antimora Apr 30, 2024
bb89dbd
Log improvements
antimora Apr 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions crates/burn-import/src/onnx/dim_inference.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,12 @@ fn cast_update_outputs(node: &mut Node) {
}
_ => panic!("Cast: only scalar and tensor inputs are valid"),
}

log::debug!(
"Cast: input type: {:?}, output type: {:?}",
input.ty,
output.ty
);
}

fn concat_update_outputs(node: &mut Node) {
Expand Down Expand Up @@ -300,11 +306,10 @@ fn same_as_input(node: &mut Node) {
}

/// Temporary pass-through stub for dimension inference so that we can export the IR model.
fn temporary_pass_through_stub(node: &Node) {
log::warn!(
"Must implement dimension inference for {:?}",
node.node_type
);
fn temporary_pass_through_stub(node: &mut Node) {
log::warn!("Must implement dimension inference for {:?}", node);
log::warn!("Temporarily setting the output type to the input type.");
node.outputs[0].ty = node.inputs[0].ty.clone();
}

fn equal_update_outputs(node: &mut Node) {
Expand Down
5 changes: 2 additions & 3 deletions crates/burn-import/src/onnx/from_onnx.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,13 +503,12 @@ fn rename_io(node: &mut Node, graph_io: &mut OnnxGraphIO) {
node_input.passed = false;
}
}
log::debug!("\n\nchecking outputs");
let mut out_count = 1;
if node.node_type == NodeType::Constant || node.node_type == NodeType::Identity {
log::debug!("it's a constant");
let new_name = format!("{}_out{}", node.name, out_count);
graph_io.insert(&node.outputs[0], &new_name);
node.outputs[0].name = new_name;
node.outputs[0].name = new_name.clone();
log::debug!("Found {} constant", new_name);
} else {
for output in node.outputs.iter_mut() {
log::debug!("output name: {}", &output.name);
Expand Down
3 changes: 2 additions & 1 deletion crates/burn-import/src/onnx/op_configuration.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -569,8 +569,9 @@ pub fn reshape_config(node: &Node) -> Vec<i64> {
panic!("Zero shape size is not supported");
}

// TODO: check "shape" attribute
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this specific to an opset that's not supported yet, hence the "TODO"?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, correct. I do not have a specific use case but since we handle similar one and I noticed from the spec, we are missing it. I just added TODO as a debt, which we do not need to make a payment on.

if node.inputs.len() != 2 || node.inputs[1].value.is_none() {
panic!("Reshape: shape tensor must be present");
panic!("Reshape: shape tensor must be present for {:?}", node);
}

let input_value = &node.inputs[1].value;
Expand Down
8 changes: 7 additions & 1 deletion crates/burn-import/src/onnx/to_burn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,8 @@ impl OnnxGraph {
pub fn into_burn<PS: PrecisionSettings + 'static>(self) -> BurnGraph<PS> {
let mut graph = BurnGraph::<PS>::default();

let mut unsupported_ops = vec![];

for node in self.nodes {
match node.node_type {
NodeType::Add => graph.register(Self::add_conversion(node)),
Expand Down Expand Up @@ -279,10 +281,14 @@ impl OnnxGraph {
NodeType::Unsqueeze => graph.register(Self::unsqueeze_conversion(node)),
NodeType::Where => graph.register(Self::where_conversion(node)),
NodeType::Sign => graph.register(Self::sign_conversion(node)),
_ => panic!("Unsupported node conversion {}", node.node_type),
node_type => unsupported_ops.push(node_type),
}
}

if !unsupported_ops.is_empty() {
panic!("Unsupported ops: {:?}", unsupported_ops);
}

// Get input and output names
let input_names = self
.inputs
Expand Down
Loading