feat: ResNet-18 end to end, Conv+BN folding, preload anonymous weights

.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

.gitignore

@@ -9,3 +9,4 @@ Cargo.lock
 .claude/
 tests/*.onnx
 tests/*.npy
+tests/models/

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.

README.md

@@ -24,9 +24,28 @@ 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, burn `flex` backend, median of 10 runs).
+Performance varies a lot by platform:
+
+| 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.
 
 ## Building
 
@@ -42,6 +61,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

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<String> {
+    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<FloatPrim> {
         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,
         }
     }

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<FloatPrim> {
-        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<Node>,
+    weights: &mut HashMap<String, FloatPrim>,
+    graph_outputs: &[String],
+) -> Vec<Node> {
+    let mut consumers: HashMap<String, usize> = 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<String, usize> = 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::<f32, _>([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<usize> = 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<OnnxModel, String> {
         }
     }
 
+    // pass 3: fold Conv -> BN pairs so those BNs cost nothing at inference time
+    let output_names: Vec<String> = 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))),

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);