From 3c2f26a45f77315d4a9875eabbed18fa39169162 Mon Sep 17 00:00:00 2001 From: Asjad Ali <133517771+asjad2401@users.noreply.github.com> Date: Mon, 5 Oct 2026 18:51:45 +0500 Subject: [PATCH 1/2] feat: ResNet-18 end to end, Conv+BN folding, preload anonymous weights - ResNet-18 now verified against ONNX Runtime on Model Zoo v1-7/v2-7 and torchvision legacy + dynamo exports (tests/test_resnet.py, new CI job). - Add ReduceMean/Sum/Max/Min; the dynamo exporter emits ReduceMean for global average pooling. - Pre-load static weights that onnx-ir passes as anonymous constants (keyed by data id). Previously every conv weight was converted from TensorData on each forward pass. - Fold BatchNormalization into a directly preceding Conv2d at load time when the conv output has no other consumer. Works around slow [1,C,1,1] broadcasting in burn-flex: Model Zoo v1 batch 8 279 -> 103 ms. - Only use the pre-fused BN fast path for 4D inputs; rank-2 BN was being broadcast to a 4D output. - Linear/Gemm bias goes through ctx.resolve (uses pre-loaded weights). Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/ci.yml | 41 +++++++++++ .gitignore | 1 + CHANGELOG.md | 12 ++++ README.md | 19 +++++- src/interpreter/context.rs | 31 +++++---- src/interpreter/mod.rs | 125 +++++++++++++++++++++++++++++----- src/interpreter/ops/conv.rs | 7 +- src/interpreter/ops/linear.rs | 46 +++---------- src/interpreter/ops/mod.rs | 1 + src/interpreter/ops/reduce.rs | 80 ++++++++++++++++++++++ tests/test_ops.py | 79 +++++++++++++++++++++ tests/test_resnet.py | 115 +++++++++++++++++++++++++++++++ 12 files changed, 489 insertions(+), 68 deletions(-) create mode 100644 src/interpreter/ops/reduce.rs create mode 100644 tests/test_resnet.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 795e010..a07abd1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,3 +84,44 @@ jobs: - name: Run op edge-case tests run: .venv/bin/python tests/test_ops.py + + resnet18-e2e: + name: resnet18 end-to-end + runs-on: ubuntu-latest + needs: build + steps: + - uses: actions/checkout@v4 + + - name: Install Rust + uses: dtolnay/rust-toolchain@stable + + - name: Cache cargo + uses: Swatinem/rust-cache@v2 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + # Model Zoo downloads, torch exports and torchvision weights (~200 MB total) + - name: Cache models + uses: actions/cache@v4 + with: + path: | + tests/models + ~/.cache/torch + key: resnet18-models-v1 + + - name: Create virtualenv + run: python -m venv .venv + + - name: Install Python deps + run: | + .venv/bin/pip install maturin numpy onnx onnxruntime onnxscript + .venv/bin/pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + + - name: Build extension + run: source .venv/bin/activate && maturin develop --release + + - name: Run ResNet-18 tests + run: .venv/bin/python tests/test_resnet.py --bench diff --git a/.gitignore b/.gitignore index 0a45f9b..b63d641 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ Cargo.lock .claude/ tests/*.onnx tests/*.npy +tests/models/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 61181a1..63a5782 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Performance +- Static weights that onnx-ir passes as anonymous (unnamed) constants — including all + conv weights — are now pre-loaded once instead of converted on every forward pass. +- BatchNormalization directly after a Conv2d is folded into the conv's weight and bias + at load time and removed from the graph. ResNet-18 (Model Zoo v1) at batch 8: 279 → 103 ms. + ### Changed - Errors now surface as normal Python exceptions instead of `PanicException`: - `TypeError` for non-array / non-float32 inputs (with a `.astype(np.float32)` hint). @@ -17,6 +23,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 set `RUST_BACKTRACE=1` to get it back. ### Fixed +- BatchNormalization on non-4D inputs (e.g. `[N, C]` after a Linear) no longer + broadcasts to a 4D output; the pre-fused fast path is only used for NCHW. - Fortran-ordered inputs (e.g. `x.T`) were read in memory order, silently scrambling data. Non-C-contiguous arrays are now copied in logical order. - Conv2d / MaxPool2d / AveragePool2d with asymmetric pads no longer average the two @@ -29,6 +37,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 being skipped with a warning and failing (or silently misbehaving) at inference time. ### Added +- ReduceMean, ReduceSum, ReduceMax, ReduceMin (needed by the torch dynamo exporter's + global average pooling). +- `tests/test_resnet.py`: ResNet-18 end to end vs ONNX Runtime on Model Zoo and + torchvision (legacy + dynamo) exports, with optional `--bench`; run in CI. - `tests/test_ops.py`: per-op edge-case tests against ONNX Runtime, run in CI. - ONNX interpreter support for Conv2d, fused BatchNormalization, MaxPool2d, AveragePool2d, and GlobalAveragePool. diff --git a/README.md b/README.md index 2a1e715..20439b2 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,23 @@ currently supports: - **Activations**: Relu, Sigmoid, Tanh, Gelu, Softmax, LogSoftmax - **Elementwise**: Add, Sub, Mul, Div - **Shape ops**: Reshape, Flatten, Transpose -- **Conv/pooling**: Conv2d, BatchNormalization (fused), MaxPool2d, AveragePool2d, GlobalAveragePool +- **Conv/pooling**: Conv2d, BatchNormalization (folded into the preceding Conv when possible), MaxPool2d, AveragePool2d, GlobalAveragePool +- **Reductions**: ReduceMean, ReduceSum, ReduceMax, ReduceMin -Enough to run simple MLPs and small CNNs; more ops are being added incrementally. +### Verified models + +ResNet-18 runs end to end and matches ONNX Runtime (max logit diff < 1e-5, same top-1) +for the ONNX Model Zoo exports (`resnet18-v1-7`, `resnet18-v2-7`) and torchvision's +`resnet18` exported with both `torch.onnx` exporters (legacy and dynamo). + +CPU latency vs ONNX Runtime (torchvision export, Apple M-series, 10 cores, median of 10): + +| batch | burn-python | ONNX Runtime | +|------:|------------:|-------------:| +| 1 | 22.6 ms | 14.7 ms | +| 8 | 105.0 ms | 118.5 ms | + +More ops are being added incrementally. ## Building @@ -42,6 +56,7 @@ python tests/make_test_model.py # generates tests/mlp.onnx python tests/test_bridge.py # numpy <-> Burn tensor bridge python tests/compare_ort.py # correctness + perf vs ONNX Runtime python tests/test_ops.py # op edge cases (padding, pooling, reshape) vs ORT +python tests/test_resnet.py # ResNet-18 end to end vs ORT (add --bench for timings) ``` ## License diff --git a/src/interpreter/context.rs b/src/interpreter/context.rs index 0f1e2ea..41435cd 100644 --- a/src/interpreter/context.rs +++ b/src/interpreter/context.rs @@ -8,6 +8,17 @@ use onnx_ir::ir::{Argument, ValueSource}; use crate::tensor::{B, FloatPrim, default_device}; +/// Key a static/constant argument is pre-loaded under in the weights map. +/// onnx-ir often hands initializers over as anonymous statics (name ""), so those +/// are keyed by their data id instead. +pub fn weight_key(arg: &Argument) -> Option { + match arg.value_source { + _ if !arg.name.is_empty() => Some(arg.name.clone()), + ValueSource::Static(id) => Some(format!("static::{id}")), + _ => None, + } +} + /// Ops report failures as a message; the interpreter adds node context. pub type OpResult = Result<(), String>; @@ -37,22 +48,18 @@ impl<'w> ExecutionContext<'w> { /// Resolve any Argument to a FloatPrim: /// - Dynamic/named Constant → look up by name - /// - Static with empty name → convert inline TensorData directly + /// - Static → pre-loaded weight (see `weight_key`), else convert inline data /// - Optional/missing → None pub fn resolve(&self, arg: &Argument) -> Option { match arg.value_source { ValueSource::Dynamic | ValueSource::Constant => self.get(&arg.name), - ValueSource::Static(_) => { - if !arg.name.is_empty() { - // pre-loaded by name - if let Some(t) = self.get(&arg.name) { - return Some(t); - } - } - // anonymous inline static — convert on the fly - arg.value() - .map(|d| B::float_from_data(d, &default_device())) - } + ValueSource::Static(_) => weight_key(arg) + .and_then(|key| self.weights.get(&key).cloned()) + // not pre-loaded (e.g. non-float data) — convert on the fly + .or_else(|| { + arg.value() + .map(|d| B::float_from_data(d, &default_device())) + }), ValueSource::Optional => None, } } diff --git a/src/interpreter/mod.rs b/src/interpreter/mod.rs index 8577f7f..6ee9675 100644 --- a/src/interpreter/mod.rs +++ b/src/interpreter/mod.rs @@ -22,22 +22,17 @@ use pyo3::{ }; use crate::tensor::{B, FloatPrim, default_device, flex_to_numpy, numpy_to_flex}; -use context::{ExecutionContext, OpResult}; +use context::{ExecutionContext, OpResult, weight_key}; fn load_weight(arg: &Argument) -> Option<(String, FloatPrim)> { - if arg.name.is_empty() { - return None; - } match arg.value_source { ValueSource::Static(_) | ValueSource::Constant => { + let key = weight_key(arg)?; let data = arg.value()?; if !matches!(data.dtype, DType::F32 | DType::F16 | DType::BF16) { return None; } - Some(( - arg.name.clone(), - B::float_from_data(data, &default_device()), - )) + Some((key, B::float_from_data(data, &default_device()))) } _ => None, } @@ -56,14 +51,7 @@ fn precompute_bn( }; // all params must be pre-loaded (static/constant) - let get = |arg: &Argument| -> Option { - if !arg.name.is_empty() { - weights.get(&arg.name).cloned() - } else { - arg.value() - .map(|d| B::float_from_data(d, &default_device())) - } - }; + let get = |arg: &Argument| weights.get(&weight_key(arg)?).cloned(); let gamma = get(&node.inputs[1])?; let beta = get(&node.inputs[2])?; @@ -86,6 +74,95 @@ fn precompute_bn( Some((node.name.clone(), scale, offset)) } +/// Fold each BatchNorm that directly follows a Conv2d into the conv's weight and bias: +/// W' = W * scale[c_out], b' = b * scale + offset +/// then drop the BN node and have the conv write the BN's output directly. +/// Only done when the conv output feeds nothing but that BN, and the BN's scale/offset +/// were pre-computed (`precompute_bn`). Unfolded BNs keep the runtime mul + add path. +fn fold_conv_bn( + mut nodes: Vec, + weights: &mut HashMap, + graph_outputs: &[String], +) -> Vec { + let mut consumers: HashMap = HashMap::new(); + for name in nodes + .iter() + .flat_map(|n| n.inputs().iter().map(|a| &a.name)) + .chain(graph_outputs) + { + *consumers.entry(name.clone()).or_default() += 1; + } + let conv_by_output: HashMap = nodes + .iter() + .enumerate() + .filter(|(_, n)| matches!(n, Node::Conv2d(_))) + .map(|(i, n)| (n.outputs()[0].name.clone(), i)) + .collect(); + + let mut folded = Vec::new(); // (conv index, bn index) + for (j, node) in nodes.iter().enumerate() { + let Node::BatchNormalization(bn) = node else { + continue; + }; + let x = &bn.inputs[0].name; + let (Some(&i), Some(1)) = (conv_by_output.get(x), consumers.get(x)) else { + continue; + }; + let Node::Conv2d(conv) = &nodes[i] else { + continue; + }; + let (Some(scale), Some(offset), Some(w)) = ( + weights.get(&format!("{}::scale", bn.name)).cloned(), + weights.get(&format!("{}::offset", bn.name)).cloned(), + weight_key(&conv.inputs[1]).and_then(|k| weights.get(&k).cloned()), + ) else { + continue; + }; + let c_out = w.shape()[0]; + let bias = match conv.inputs.get(2) { + Some(arg) => match weight_key(arg).and_then(|k| weights.get(&k).cloned()) { + Some(b) => b, + None => continue, // dynamic bias: leave unfolded + }, + None => B::float_from_data(TensorData::zeros::([c_out]), &default_device()), + }; + + let w = B::float_mul( + w, + B::float_reshape(scale.clone(), vec![c_out, 1, 1, 1].into()), + ); + let b = B::float_add( + B::float_mul(bias, B::float_reshape(scale, vec![c_out].into())), + B::float_reshape(offset, vec![c_out].into()), + ); + weights.insert(format!("{}::folded_weight", conv.name), w); + weights.insert(format!("{}::folded_bias", conv.name), b); + folded.push((i, j)); + } + + for &(i, j) in &folded { + let bn_out = nodes[j].outputs()[0].name.clone(); + let Node::Conv2d(conv) = &mut nodes[i] else { + unreachable!() + }; + // point the conv at the folded params (static args are looked up by name first) + let mut w_arg = conv.inputs[1].clone(); + w_arg.name = format!("{}::folded_weight", conv.name); + let mut b_arg = w_arg.clone(); + b_arg.name = format!("{}::folded_bias", conv.name); + conv.inputs.truncate(1); + conv.inputs.extend([w_arg, b_arg]); + conv.outputs[0].name = bn_out; + } + let dropped: std::collections::HashSet = folded.iter().map(|&(_, j)| j).collect(); + nodes + .into_iter() + .enumerate() + .filter(|(j, _)| !dropped.contains(j)) + .map(|(_, n)| n) + .collect() +} + /// A graph input, with whatever shape info the model declares (None = symbolic dim). struct InputSpec { name: String, @@ -274,10 +351,14 @@ pub fn load_onnx(path: &str) -> Result { } } + // pass 3: fold Conv -> BN pairs so those BNs cost nothing at inference time + let output_names: Vec = graph.outputs.iter().map(|a| a.name.clone()).collect(); + let nodes = fold_conv_bn(graph.nodes, &mut weights, &output_names); + Ok(OnnxModel { inputs: graph.inputs.iter().map(InputSpec::from_arg).collect(), - output_names: graph.outputs.iter().map(|a| a.name.clone()).collect(), - nodes: graph.nodes, + output_names, + nodes, weights, }) } @@ -306,6 +387,10 @@ fn is_supported(node: &Node) -> bool { | Node::MaxPool2d(_) | Node::AveragePool2d(_) | Node::GlobalAveragePool(_) + | Node::ReduceMean(_) + | Node::ReduceSum(_) + | Node::ReduceMax(_) + | Node::ReduceMin(_) | Node::Constant(_) ) } @@ -338,6 +423,10 @@ fn dispatch(node: &Node, ctx: &mut ExecutionContext) -> OpResult { Node::MaxPool2d(n) => ops::pool::max_pool2d(n, ctx), Node::AveragePool2d(n) => ops::pool::avg_pool2d(n, ctx), Node::GlobalAveragePool(n) => ops::pool::global_avg_pool(n, ctx), + Node::ReduceMean(n) => ops::reduce::mean(n, ctx), + Node::ReduceSum(n) => ops::reduce::sum(n, ctx), + Node::ReduceMax(n) => ops::reduce::max(n, ctx), + Node::ReduceMin(n) => ops::reduce::min(n, ctx), Node::Constant(_) => Ok(()), // rejected by `is_supported` at load time other => Err(format!("unsupported op '{}'", op_type(other))), diff --git a/src/interpreter/ops/conv.rs b/src/interpreter/ops/conv.rs index e7bb61b..2999b0b 100644 --- a/src/interpreter/ops/conv.rs +++ b/src/interpreter/ops/conv.rs @@ -50,7 +50,12 @@ pub fn batch_norm_fused(node: &BatchNormalizationNode, ctx: &mut ExecutionContex let scale_key = format!("{}::scale", node.name); let offset_key = format!("{}::offset", node.name); - if let (Some(scale), Some(offset)) = (ctx.get(&scale_key), ctx.get(&offset_key)) { + // pre-fused scale/offset are shaped [1, C, 1, 1], so only valid for NCHW inputs + let fused = match x.shape().num_dims() { + 4 => ctx.get(&scale_key).zip(ctx.get(&offset_key)), + _ => None, + }; + if let Some((scale, offset)) = fused { // fast path: 2 ops, no reshape (scale/offset are already [1,C,1,1]) let y = B::float_mul(x, scale); let y = B::float_add(y, offset); diff --git a/src/interpreter/ops/linear.rs b/src/interpreter/ops/linear.rs index 813a642..73fb52b 100644 --- a/src/interpreter/ops/linear.rs +++ b/src/interpreter/ops/linear.rs @@ -18,23 +18,8 @@ pub fn linear(node: &LinearNode, ctx: &mut ExecutionContext) -> OpResult { let y = B::float_matmul(x, w); // bias is optional (3rd input) - let y = match node.inputs.get(2) { - Some(bias_arg) if !bias_arg.name.is_empty() => { - if let Some(b) = ctx.get(&bias_arg.name) { - B::float_add(y, b) - } else { - y - } - } - Some(bias_arg) => { - // inline static (name="") - if let Some(data) = bias_arg.value() { - let b = B::float_from_data(data, &default_device()); - B::float_add(y, b) - } else { - y - } - } + let y = match node.inputs.get(2).and_then(|a| ctx.resolve(a)) { + Some(b) => B::float_add(y, b), None => y, }; @@ -69,26 +54,17 @@ pub fn gemm(node: &GemmNode, ctx: &mut ExecutionContext) -> OpResult { y = B::float_mul(y, alpha); } - if let Some(c_arg) = node.inputs.get(2) { - let c = if !c_arg.name.is_empty() { - ctx.get(&c_arg.name) + if let Some(c) = node.inputs.get(2).and_then(|a| ctx.resolve(a)) { + let c = if node.config.beta != 1.0 { + let beta = B::float_from_data( + burn_backend::TensorData::from([node.config.beta]), + &default_device(), + ); + B::float_mul(c, beta) } else { - c_arg - .value() - .map(|d| B::float_from_data(d, &default_device())) + c }; - if let Some(c) = c { - let c = if node.config.beta != 1.0 { - let beta = B::float_from_data( - burn_backend::TensorData::from([node.config.beta]), - &default_device(), - ); - B::float_mul(c, beta) - } else { - c - }; - y = B::float_add(y, c); - } + y = B::float_add(y, c); } ctx.insert(node.outputs[0].name.clone(), y); diff --git a/src/interpreter/ops/mod.rs b/src/interpreter/ops/mod.rs index 7e6a1ba..0e8e60e 100644 --- a/src/interpreter/ops/mod.rs +++ b/src/interpreter/ops/mod.rs @@ -4,4 +4,5 @@ pub mod elementwise; pub mod linear; mod padding; pub mod pool; +pub mod reduce; pub mod reshape; diff --git a/src/interpreter/ops/reduce.rs b/src/interpreter/ops/reduce.rs new file mode 100644 index 0000000..77ff0c3 --- /dev/null +++ b/src/interpreter/ops/reduce.rs @@ -0,0 +1,80 @@ +use super::super::context::{ExecutionContext, OpResult}; +use crate::tensor::{B, FloatPrim}; +use burn_backend::{TensorMetadata, backend::ops::FloatTensorOps}; +use onnx_ir::{ + ir::Argument, + reduce::{ReduceConfig, ReduceMaxNode, ReduceMeanNode, ReduceMinNode, ReduceSumNode}, +}; + +// onnx-ir hands us sorted, non-negative dims; empty means "reduce everything". +// Burn's *_dim ops keep the reduced dim as size 1, matching keepdims=1. +fn reduce( + inputs: &[Argument], + outputs: &[Argument], + cfg: &ReduceConfig, + ctx: &mut ExecutionContext, + op: fn(FloatPrim, usize) -> FloatPrim, +) -> OpResult { + let x = ctx.require(&inputs[0])?; + let rank = x.shape().num_dims(); + let dims: Vec = if cfg.dims.is_empty() { + (0..rank).collect() + } else { + cfg.dims.clone() + }; + + let mut y = dims.iter().fold(x, |t, &d| op(t, d)); + if !cfg.keepdims { + let shape: Vec = y + .shape() + .iter() + .enumerate() + .filter(|(i, _)| !dims.contains(i)) + .map(|(_, &s)| s) + .collect(); + y = B::float_reshape(y, shape.into()); + } + + ctx.insert(outputs[0].name.clone(), y); + Ok(()) +} + +pub fn mean(node: &ReduceMeanNode, ctx: &mut ExecutionContext) -> OpResult { + reduce( + &node.inputs, + &node.outputs, + &node.config, + ctx, + B::float_mean_dim, + ) +} + +pub fn sum(node: &ReduceSumNode, ctx: &mut ExecutionContext) -> OpResult { + reduce( + &node.inputs, + &node.outputs, + &node.config, + ctx, + B::float_sum_dim, + ) +} + +pub fn max(node: &ReduceMaxNode, ctx: &mut ExecutionContext) -> OpResult { + reduce( + &node.inputs, + &node.outputs, + &node.config, + ctx, + B::float_max_dim, + ) +} + +pub fn min(node: &ReduceMinNode, ctx: &mut ExecutionContext) -> OpResult { + reduce( + &node.inputs, + &node.outputs, + &node.config, + ctx, + B::float_min_dim, + ) +} diff --git a/tests/test_ops.py b/tests/test_ops.py index 00d83a0..3befd61 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -91,6 +91,69 @@ def pool(op, **attrs): compare("avgpool SAME_LOWER s2", pool("AveragePool", auto_pad="SAME_LOWER", strides=[2, 2]), X) +# ── batchnorm (folded into conv where possible) ─────────────────────────────── + +print("\n=== batchnorm ===") + +def graph_model(name, nodes, x_shape, initializers, outputs=("y",)): + graph = helper.make_graph( + nodes, name, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, list(x_shape))], + [helper.make_tensor_value_info(o, TensorProto.FLOAT, None) for o in outputs], + initializer=list(initializers), + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 10 + path = os.path.join(tmpdir, f"{name}.onnx") + onnx.save(model, path) + return path + +def compare_graph(label, nodes, x_shape, initializers, outputs=("y",), n_nodes=None): + path = graph_model(label.replace(" ", "_"), nodes, x_shape, initializers, outputs) + x = np.random.randn(*x_shape).astype(np.float32) + try: + model = burn.load_onnx(path) + if n_nodes is not None: # e.g. BN folded away into the conv + check(f"{label} runs {n_nodes} node(s)", f"nodes={n_nodes})" in repr(model)) + got = model([x]) + except BaseException as e: + check(f"{label} ({type(e).__name__}: {e})", False) + return + want = ort.InferenceSession(path, providers=["CPUExecutionProvider"]).run(None, {"x": x}) + ok = all(g.shape == w.shape and np.allclose(g, w, atol=1e-4) for g, w in zip(got, want)) + diff = max(np.abs(g - w).max() for g, w in zip(got, want)) if ok else float("nan") + check(f"{label} max_abs_diff={diff:.2e}", ok) + +def bn_params(c, prefix="bn"): + return [ + numpy_helper.from_array(np.random.rand(c).astype(np.float32) + 0.5, f"{prefix}_g"), + numpy_helper.from_array(np.random.randn(c).astype(np.float32), f"{prefix}_b"), + numpy_helper.from_array(np.random.randn(c).astype(np.float32), f"{prefix}_m"), + numpy_helper.from_array(np.random.rand(c).astype(np.float32) + 0.5, f"{prefix}_v"), + ] + +def bn(x, y, prefix="bn"): + return helper.make_node("BatchNormalization", + [x, f"{prefix}_g", f"{prefix}_b", f"{prefix}_m", f"{prefix}_v"], [y]) + +compare_graph("conv+bn (folded)", + [helper.make_node("Conv", ["x", "W", "B"], ["c"], kernel_shape=[3, 3], pads=[1, 1, 1, 1]), + bn("c", "y")], X, [W, B] + bn_params(3), n_nodes=1) +compare_graph("conv(no bias)+bn (folded)", + [helper.make_node("Conv", ["x", "W"], ["c"], kernel_shape=[3, 3]), + bn("c", "y")], X, [W] + bn_params(3), n_nodes=1) +compare_graph("conv+bn, conv output reused (not folded)", + [helper.make_node("Conv", ["x", "W", "B"], ["c"], kernel_shape=[3, 3]), + bn("c", "n"), helper.make_node("Add", ["c", "n"], ["y"])], X, [W, B] + bn_params(3), + n_nodes=3) +compare_graph("conv+bn, conv output is graph output (not folded)", + [helper.make_node("Conv", ["x", "W", "B"], ["c"], kernel_shape=[3, 3]), + bn("c", "y")], X, [W, B] + bn_params(3), outputs=("y", "c"), n_nodes=2) +compare_graph("bn after relu (runtime path)", + [helper.make_node("Relu", ["x"], ["r"]), bn("r", "y", "bn2")], X, bn_params(2, "bn2")) +compare_graph("bn on rank-2 input", + [bn("x", "y", "bn3")], (5, 4), bn_params(4, "bn3")) + # ── reshape ────────────────────────────────────────────────────────────────── print("\n=== reshape ===") @@ -106,6 +169,22 @@ def reshape(shape): node, init = reshape([-1, 4]) compare("reshape [-1, 4]", node, (2, 3, 4), init) +# ── reductions ─────────────────────────────────────────────────────────────── + +print("\n=== reductions ===") + +def reduce_node(op, axes, keepdims): + a = numpy_helper.from_array(np.array(axes, dtype=np.int64), "axes") + return helper.make_node(op, ["x", "axes"], ["y"], keepdims=keepdims), [a] + +for op in ["ReduceMean", "ReduceSum", "ReduceMax", "ReduceMin"]: + node, init = reduce_node(op, [-1, -2], 1) + compare(f"{op} axes=[-1,-2] keepdims=1", node, X, init, opset=18) + node, init = reduce_node(op, [1], 0) + compare(f"{op} axes=[1] keepdims=0", node, X, init, opset=18) +compare("ReduceMean all axes keepdims=1", + helper.make_node("ReduceMean", ["x"], ["y"], keepdims=1), X, opset=18) + # ── load errors ────────────────────────────────────────────────────────────── print("\n=== load errors ===") diff --git a/tests/test_resnet.py b/tests/test_resnet.py new file mode 100644 index 0000000..146952b --- /dev/null +++ b/tests/test_resnet.py @@ -0,0 +1,115 @@ +""" +End-to-end ResNet-18 vs ONNX Runtime, on real exported models: + - ONNX Model Zoo resnet18-v1-7 / v2-7 (MXNet exports, downloaded + cached) + - torchvision resnet18 via torch.onnx legacy + dynamo exporters (if torch is installed) + +Checks logits match ORT and top-1 agrees, across batch sizes. +Run: python tests/test_resnet.py [--bench] +""" + +import os +import sys +import time +import urllib.request +import numpy as np +import onnxruntime as ort +import burn_python as burn + +MODEL_DIR = os.path.join(os.path.dirname(__file__), "models") +ZOO_URL = "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/{}.onnx" +BATCHES = [1, 4] +BENCH = "--bench" in sys.argv + +PASS = "\033[92mPASS\033[0m" +FAIL = "\033[91mFAIL\033[0m" + +failures = 0 + +def check(label, cond): + global failures + print(f" {'[' + PASS + ']' if cond else '[' + FAIL + ']'} {label}") + if not cond: + failures += 1 + return cond + +def zoo_model(name): + path = os.path.join(MODEL_DIR, f"{name}.onnx") + if not os.path.exists(path): + print(f"downloading {name} ...") + urllib.request.urlretrieve(ZOO_URL.format(name), path + ".part") + os.rename(path + ".part", path) + return path + +def torch_models(): + try: + import torch + import torchvision + except ImportError: + print("torch/torchvision not installed — skipping PyTorch exports") + return [] + + model = torchvision.models.resnet18(weights="IMAGENET1K_V1").eval() + # batch 2, not 1: torch.export specializes size-1 dims, baking batch=1 into Reshape + x = torch.randn(2, 3, 224, 224) + paths = [] + + legacy = os.path.join(MODEL_DIR, "resnet18-torch-legacy.onnx") + if not os.path.exists(legacy): + torch.onnx.export(model, x, legacy, input_names=["input"], output_names=["logits"], + dynamic_axes={"input": {0: "N"}, "logits": {0: "N"}}, + opset_version=17, dynamo=False) + paths.append(legacy) + + dynamo = os.path.join(MODEL_DIR, "resnet18-torch-dynamo.onnx") + if not os.path.exists(dynamo): + try: + torch.onnx.export(model, (x,), dynamo, input_names=["input"], output_names=["logits"], + dynamic_shapes={"x": {0: torch.export.Dim("N")}}, dynamo=True) + except Exception as e: # dynamo exporter needs onnxscript + print(f"dynamo export unavailable ({type(e).__name__}) — skipping") + return paths + paths.append(dynamo) + return paths + +def median_ms(fn, runs=10): + times = [] + for _ in range(runs): + t0 = time.perf_counter() + fn() + times.append(time.perf_counter() - t0) + return np.median(times) * 1000 + +os.makedirs(MODEL_DIR, exist_ok=True) +models = [zoo_model("resnet18-v1-7"), zoo_model("resnet18-v2-7")] + torch_models() + +for path in models: + print(f"\n=== {os.path.basename(path)} ===") + model = burn.load_onnx(path) + sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"]) + input_name = sess.get_inputs()[0].name + + rng = np.random.RandomState(0) + for batch in BATCHES: + x = rng.randn(batch, 3, 224, 224).astype(np.float32) + got = model([x])[0] + want = sess.run(None, {input_name: x})[0] + diff = np.abs(got - want).max() + check(f"batch={batch} max_abs_diff={diff:.2e}", + got.shape == want.shape and np.allclose(got, want, atol=1e-4, rtol=1e-4)) + check(f"batch={batch} top-1 matches ORT", np.array_equal(got.argmax(1), want.argmax(1))) + + if BENCH: + for batch in (1, 8): + x = rng.randn(batch, 3, 224, 224).astype(np.float32) + for _ in range(3): # warm up + model([x]) + sess.run(None, {input_name: x}) + burn_ms = median_ms(lambda: model([x])) + ort_ms = median_ms(lambda: sess.run(None, {input_name: x})) + print(f" bench batch={batch:<2} burn {burn_ms:6.1f} ms | ORT {ort_ms:6.1f} ms | ratio {burn_ms / ort_ms:.2f}x") + +print() + +if failures: + print(f"{failures} check(s) failed") + sys.exit(1) From c9627fab011435da633e8600608bff8decdb0a81 Mon Sep 17 00:00:00 2001 From: Asjad Ali <133517771+asjad2401@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:28:27 +0500 Subject: [PATCH 2/2] docs: add Linux x86 ResNet-18 benchmark numbers alongside Apple M-series Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 20439b2..666201e 100644 --- a/README.md +++ b/README.md @@ -33,12 +33,17 @@ ResNet-18 runs end to end and matches ONNX Runtime (max logit diff < 1e-5, same for the ONNX Model Zoo exports (`resnet18-v1-7`, `resnet18-v2-7`) and torchvision's `resnet18` exported with both `torch.onnx` exporters (legacy and dynamo). -CPU latency vs ONNX Runtime (torchvision export, Apple M-series, 10 cores, median of 10): +CPU latency vs ONNX Runtime (torchvision export, burn `flex` backend, median of 10 runs). +Performance varies a lot by platform: -| batch | burn-python | ONNX Runtime | -|------:|------------:|-------------:| -| 1 | 22.6 ms | 14.7 ms | -| 8 | 105.0 ms | 118.5 ms | +| platform | batch | burn-python | ONNX Runtime | burn / ORT | +|----------|------:|------------:|-------------:|-----------:| +| Apple M-series (10 cores) | 1 | 22.6 ms | 14.7 ms | 1.5× slower | +| Apple M-series (10 cores) | 8 | 105.0 ms | 118.5 ms | 0.9× (faster) | +| Linux x86 (GitHub Actions runner) | 1 | 58.0 ms | 18.7 ms | 3.1× slower | +| Linux x86 (GitHub Actions runner) | 8 | 335.7 ms | 145.9 ms | 2.3× slower | + +Run `python tests/test_resnet.py --bench` to measure on your machine. More ops are being added incrementally.