diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2c98001..372052f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,8 +3,8 @@ name: CI on: push: branches: [main] + # all PRs, so stacked PRs (based on another branch) get checked too pull_request: - branches: [main] env: CARGO_TERM_COLOR: always @@ -81,3 +81,6 @@ jobs: - name: Run ORT comparison tests run: .venv/bin/python tests/compare_ort.py + + - name: Run op edge-case tests + run: .venv/bin/python tests/test_ops.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0aead5f..d92591f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- Conv2d / MaxPool2d / AveragePool2d with asymmetric pads no longer average the two + sides; uneven pads are applied to the input explicitly (exact for `count_include_pad=0`). +- `auto_pad=SAME_UPPER` / `SAME_LOWER` / `VALID` are now honoured instead of treated as no padding. +- Pooling now respects `ceil_mode`. +- Reshape with both `0` and `-1` in the target shape no longer panics (divide by zero), + and the opset < 5 `shape` attribute form is supported. +- Unsupported ops now fail at `load_onnx` with a list of the missing ops, instead of + being skipped with a warning and failing (or silently misbehaving) at inference time. + ### Added +- `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. - GitHub Actions CI: `cargo build`, `cargo fmt` / `cargo clippy` (non-blocking), diff --git a/README.md b/README.md index 6be935a..2a1e715 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,7 @@ maturin develop --release 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 ``` ## License diff --git a/src/interpreter/mod.rs b/src/interpreter/mod.rs index c1de952..4b31cab 100644 --- a/src/interpreter/mod.rs +++ b/src/interpreter/mod.rs @@ -138,6 +138,22 @@ pub fn load_onnx(path: &str) -> Result { .parse_file(path) .map_err(|e| e.to_string())?; + // fail fast: a skipped op would only surface later as a missing tensor or wrong output + let mut unsupported: Vec = graph + .nodes + .iter() + .filter(|n| !is_supported(n)) + .map(op_type) + .collect(); + if !unsupported.is_empty() { + unsupported.sort(); + unsupported.dedup(); + return Err(format!( + "unsupported ONNX op(s): {}", + unsupported.join(", ") + )); + } + let mut weights = HashMap::new(); // pass 1: load all named static weights @@ -167,6 +183,40 @@ pub fn load_onnx(path: &str) -> Result { }) } +// keep in sync with `dispatch` +fn is_supported(node: &Node) -> bool { + matches!( + node, + Node::Relu(_) + | Node::Sigmoid(_) + | Node::Tanh(_) + | Node::Gelu(_) + | Node::Softmax(_) + | Node::LogSoftmax(_) + | Node::Linear(_) + | Node::Gemm(_) + | Node::Add(_) + | Node::Sub(_) + | Node::Mul(_) + | Node::Div(_) + | Node::Reshape(_) + | Node::Flatten(_) + | Node::Transpose(_) + | Node::Conv2d(_) + | Node::BatchNormalization(_) + | Node::MaxPool2d(_) + | Node::AveragePool2d(_) + | Node::GlobalAveragePool(_) + | Node::Constant(_) + ) +} + +/// Op type of a node (the enum variant name, e.g. "MatMul"), for error messages. +fn op_type(node: &Node) -> String { + let dbg = format!("{node:?}"); + dbg.split('(').next().unwrap_or(&dbg).to_string() +} + fn dispatch(node: &Node, ctx: &mut ExecutionContext) { match node { Node::Relu(n) => ops::activation::relu(n, ctx), @@ -190,8 +240,7 @@ fn dispatch(node: &Node, ctx: &mut ExecutionContext) { Node::AveragePool2d(n) => ops::pool::avg_pool2d(n, ctx), Node::GlobalAveragePool(n) => ops::pool::global_avg_pool(n, ctx), Node::Constant(_) => {} - other => { - eprintln!("warn: unimplemented op '{}' — skipping", other.name()); - } + // rejected by `is_supported` at load time + other => unreachable!("unsupported op '{}'", op_type(other)), } } diff --git a/src/interpreter/ops/conv.rs b/src/interpreter/ops/conv.rs index 72c859e..3444a75 100644 --- a/src/interpreter/ops/conv.rs +++ b/src/interpreter/ops/conv.rs @@ -1,4 +1,5 @@ use super::super::context::ExecutionContext; +use super::padding::{pad2d, resolve_pads, symmetric}; use crate::tensor::{B, default_device}; use burn_backend::{ TensorData, TensorMetadata, @@ -8,18 +9,8 @@ use burn_backend::{ use onnx_ir::{ batch_norm::{BatchNormConfig, BatchNormalizationNode}, conv2d::Conv2dNode, - node::padding::PaddingConfig2d, }; -fn symmetric_padding(cfg: &PaddingConfig2d) -> [usize; 2] { - match cfg { - PaddingConfig2d::Valid => [0, 0], - PaddingConfig2d::Explicit(top, left, bottom, right) => { - [(top + bottom) / 2, (left + right) / 2] - } - } -} - pub fn conv2d(node: &Conv2dNode, ctx: &mut ExecutionContext) { let x = ctx.resolve(&node.inputs[0]).expect("conv2d: missing input"); let w = ctx @@ -27,7 +18,21 @@ pub fn conv2d(node: &Conv2dNode, ctx: &mut ExecutionContext) { .expect("conv2d: missing weight"); let bias = node.inputs.get(2).and_then(|a| ctx.resolve(a)); - let pad = symmetric_padding(&node.config.padding); + let cfg = &node.config; + let pads = resolve_pads( + &x, + &cfg.padding, + &cfg.auto_pad, + cfg.kernel_size, + cfg.stride, + cfg.dilation, + ); + // uneven pads: zero-pad the input up front and run the conv unpadded + let (x, pad) = match symmetric(pads) { + Some(pad) => (x, pad), + None => (pad2d(x, pads, 0.0), [0, 0]), + }; + let options = ConvOptions::new( node.config.stride, pad, diff --git a/src/interpreter/ops/mod.rs b/src/interpreter/ops/mod.rs index 7e098b5..7e6a1ba 100644 --- a/src/interpreter/ops/mod.rs +++ b/src/interpreter/ops/mod.rs @@ -2,5 +2,6 @@ pub mod activation; pub mod conv; pub mod elementwise; pub mod linear; +mod padding; pub mod pool; pub mod reshape; diff --git a/src/interpreter/ops/padding.rs b/src/interpreter/ops/padding.rs new file mode 100644 index 0000000..c341417 --- /dev/null +++ b/src/interpreter/ops/padding.rs @@ -0,0 +1,72 @@ +// Spatial padding shared by conv and pooling ops. +// +// Burn's conv/pool kernels only take symmetric padding ([h, w] applied to both sides). +// ONNX allows asymmetric pads and auto_pad=SAME_*, so we resolve the real per-side +// pads here and, when they're uneven, pad the input explicitly instead. + +use burn_backend::{Slice, TensorData, TensorMetadata, backend::ops::FloatTensorOps}; +use onnx_ir::node::padding::{AutoPad, PaddingConfig2d}; + +use crate::tensor::{B, FloatPrim, default_device}; + +/// Per-side pads for an NCHW tensor: [top, left, bottom, right]. +pub type Pads2d = [usize; 4]; + +/// Resolve the actual pads for a 2D window op, honouring auto_pad. +pub fn resolve_pads( + x: &FloatPrim, + padding: &PaddingConfig2d, + auto_pad: &AutoPad, + kernel: [usize; 2], + stride: [usize; 2], + dilation: [usize; 2], +) -> Pads2d { + match auto_pad { + AutoPad::Valid => [0; 4], + AutoPad::SameUpper | AutoPad::SameLower => { + let shape = x.shape(); + let mut begin = [0; 2]; + let mut end = [0; 2]; + for i in 0..2 { + // output = ceil(input / stride); pad just enough to produce it + let input = shape[2 + i]; + let out = input.div_ceil(stride[i]); + let span = (kernel[i] - 1) * dilation[i] + 1; + let total = ((out - 1) * stride[i] + span).saturating_sub(input); + // SAME_UPPER puts the extra pixel at the end, SAME_LOWER at the start + let small = total / 2; + (begin[i], end[i]) = match auto_pad { + AutoPad::SameUpper => (small, total - small), + _ => (total - small, small), + }; + } + [begin[0], begin[1], end[0], end[1]] + } + AutoPad::NotSet => match padding { + PaddingConfig2d::Valid => [0; 4], + PaddingConfig2d::Explicit(top, left, bottom, right) => [*top, *left, *bottom, *right], + }, + } +} + +/// Some([h, w]) if the pads can be passed straight to a Burn kernel. +pub fn symmetric(pads: Pads2d) -> Option<[usize; 2]> { + let [top, left, bottom, right] = pads; + (top == bottom && left == right).then_some([top, left]) +} + +/// Pad an NCHW tensor's spatial dims with a constant value. +pub fn pad2d(x: FloatPrim, pads: Pads2d, value: f32) -> FloatPrim { + let [top, left, bottom, right] = pads; + let shape = x.shape(); + let (n, c, h, w) = (shape[0], shape[1], shape[2], shape[3]); + + let padded = B::float_from_data( + TensorData::full(vec![n, c, h + top + bottom, w + left + right], value), + &default_device(), + ); + let range = + |start: usize, len: usize| Slice::new(start as isize, Some((start + len) as isize), 1); + let slices = [range(0, n), range(0, c), range(top, h), range(left, w)]; + B::float_slice_assign(padded, &slices, x) +} diff --git a/src/interpreter/ops/pool.rs b/src/interpreter/ops/pool.rs index 38865e6..27f46ad 100644 --- a/src/interpreter/ops/pool.rs +++ b/src/interpreter/ops/pool.rs @@ -1,32 +1,40 @@ use super::super::context::ExecutionContext; -use crate::tensor::B; -use burn_backend::backend::ops::ModuleOps; +use super::padding::{pad2d, resolve_pads, symmetric}; +use crate::tensor::{B, default_device}; +use burn_backend::{ + TensorData, TensorMetadata, + backend::ops::{FloatTensorOps, ModuleOps}, +}; use onnx_ir::{ avg_pool2d::AveragePool2dNode, global_avg_pool::GlobalAveragePoolNode, - max_pool2d::MaxPool2dNode, node::padding::PaddingConfig2d, + max_pool2d::MaxPool2dNode, }; -fn symmetric_padding(cfg: &PaddingConfig2d) -> [usize; 2] { - match cfg { - PaddingConfig2d::Valid => [0, 0], - PaddingConfig2d::Explicit(top, left, bottom, right) => { - [(top + bottom) / 2, (left + right) / 2] - } - } -} - pub fn max_pool2d(node: &MaxPool2dNode, ctx: &mut ExecutionContext) { let x = ctx .resolve(&node.inputs[0]) .expect("max_pool2d: missing input"); - let pad = symmetric_padding(&node.config.padding); + let cfg = &node.config; + let pads = resolve_pads( + &x, + &cfg.padding, + &cfg.auto_pad, + cfg.kernel_size, + cfg.strides, + cfg.dilation, + ); + // uneven pads: -inf padding never wins the max, so pre-padding is exact + let (x, pad) = match symmetric(pads) { + Some(pad) => (x, pad), + None => (pad2d(x, pads, f32::NEG_INFINITY), [0, 0]), + }; let y = B::max_pool2d( x, - node.config.kernel_size, - node.config.strides, + cfg.kernel_size, + cfg.strides, pad, - node.config.dilation, - false, + cfg.dilation, + cfg.ceil_mode, ); ctx.insert(node.outputs[0].name.clone(), y); } @@ -35,15 +43,42 @@ pub fn avg_pool2d(node: &AveragePool2dNode, ctx: &mut ExecutionContext) { let x = ctx .resolve(&node.inputs[0]) .expect("avg_pool2d: missing input"); - let pad = symmetric_padding(&node.config.padding); - let y = B::avg_pool2d( - x, - node.config.kernel_size, - node.config.strides, - pad, - node.config.count_include_pad, - false, + let cfg = &node.config; + let pads = resolve_pads( + &x, + &cfg.padding, + &cfg.auto_pad, + cfg.kernel_size, + cfg.strides, + cfg.dilation, ); + let pool = |t, pad, include_pad| { + B::avg_pool2d( + t, + cfg.kernel_size, + cfg.strides, + pad, + include_pad, + cfg.ceil_mode, + ) + }; + + let y = match symmetric(pads) { + Some(pad) => pool(x, pad, cfg.count_include_pad), + None if cfg.count_include_pad => pool(pad2d(x, pads, 0.0), [0, 0], true), + None => { + // uneven pads that must not count towards the average: + // sum(window) / count(real pixels in window), via a padded ones-mask + let shape = x.shape(); + let ones = B::float_from_data( + TensorData::full(vec![1, 1, shape[2], shape[3]], 1.0f32), + &default_device(), + ); + let sum = pool(pad2d(x, pads, 0.0), [0, 0], true); + let count = pool(pad2d(ones, pads, 0.0), [0, 0], true); + B::float_div(sum, count) + } + }; ctx.insert(node.outputs[0].name.clone(), y); } diff --git a/src/interpreter/ops/reshape.rs b/src/interpreter/ops/reshape.rs index 1b92d01..ae9b0b3 100644 --- a/src/interpreter/ops/reshape.rs +++ b/src/interpreter/ops/reshape.rs @@ -1,39 +1,43 @@ use super::super::context::ExecutionContext; use crate::tensor::B; use burn_backend::{TensorMetadata, backend::ops::FloatTensorOps}; -use onnx_ir::{flatten::FlattenNode, reshape::ReshapeNode, transpose::TransposeNode}; +use onnx_ir::{ + flatten::FlattenNode, + reshape::{ReshapeInput, ReshapeNode}, + transpose::TransposeNode, +}; pub fn reshape(node: &ReshapeNode, ctx: &mut ExecutionContext) { let x = ctx .get(&node.inputs[0].name) .expect("reshape: missing input"); let orig_shape: Vec = x.shape().iter().copied().collect(); - let total: usize = orig_shape.iter().product(); - let shape_data = node.inputs[1] - .value() - .expect("reshape: shape must be static"); - let raw: Vec = shape_data.to_vec().expect("reshape: shape must be i64"); + // shape comes from the attribute (opset < 5) or a constant input (opset 5+) + let raw: Vec = match &node.config.shape { + ReshapeInput::Static(shape) => shape.clone(), + ReshapeInput::Runtime(r) => node.inputs[r.input_index] + .value() + .expect("reshape: shape must be static") + .to_vec() + .expect("reshape: shape must be i64"), + }; - let shape: Vec = raw + // 0 copies the input dim (allowzero=0); -1 is inferred from what's left + let mut shape: Vec = raw .iter() .enumerate() - .map(|(i, &s)| { - if s == -1 { - let known: usize = raw - .iter() - .enumerate() - .filter(|(j, v)| *j != i && **v != -1) - .map(|(_, &v)| v as usize) - .product(); - total / known - } else if s == 0 { - orig_shape[i] - } else { - s as usize - } + .map(|(i, &s)| match s { + 0 => orig_shape[i], + -1 => 1, + s => s as usize, }) .collect(); + if let Some(i) = raw.iter().position(|&s| s == -1) { + let total: usize = orig_shape.iter().product(); + let known: usize = shape.iter().product(); + shape[i] = total / known; + } ctx.insert( node.outputs[0].name.clone(), diff --git a/tests/make_test_model.py b/tests/make_test_model.py index b69bff4..44b858b 100644 --- a/tests/make_test_model.py +++ b/tests/make_test_model.py @@ -33,6 +33,8 @@ ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) +# pin the IR version: new onnx releases can write IR versions onnxruntime can't read yet +model.ir_version = 10 onnx.save(model, "tests/mlp.onnx") print("saved tests/mlp.onnx") diff --git a/tests/test_ops.py b/tests/test_ops.py new file mode 100644 index 0000000..5331674 --- /dev/null +++ b/tests/test_ops.py @@ -0,0 +1,124 @@ +""" +Op edge-case tests vs ONNX Runtime — padding modes, pooling, reshape, load errors. +Each case builds a single-node ONNX model on the fly. +Run with: python tests/test_ops.py +""" + +import os +import sys +import tempfile +import numpy as np +import onnx +import onnxruntime as ort +from onnx import helper, numpy_helper, TensorProto +import burn_python as burn + +PASS = "\033[92mPASS\033[0m" +FAIL = "\033[91mFAIL\033[0m" + +failures = 0 +tmpdir = tempfile.mkdtemp() + +def check(label, cond): + global failures + print(f" {'[' + PASS + ']' if cond else '[' + FAIL + ']'} {label}") + if not cond: + failures += 1 + return cond + +def save_model(name, node, x_shape, initializers=(), opset=17): + graph = helper.make_graph( + [node], + name, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, list(x_shape))], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + initializer=list(initializers), + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + model.ir_version = 10 # older onnxruntime builds reject newer IR versions + path = os.path.join(tmpdir, f"{name}.onnx") + onnx.save(model, path) + return path + +def compare(label, node, x_shape, initializers=(), opset=17): + path = save_model(label.replace(" ", "_"), node, x_shape, initializers, opset) + x = np.random.randn(*x_shape).astype(np.float32) + try: + got = burn.load_onnx(path)([x])[0] + except BaseException as e: # pyo3 panics aren't Exception subclasses + check(f"{label} ({type(e).__name__}: {e})", False) + return + want = ort.InferenceSession(path, providers=["CPUExecutionProvider"]).run(None, {"x": x})[0] + if got.shape != want.shape: + check(f"{label} shape {got.shape} != ORT {want.shape}", False) + return + diff = np.abs(got - want).max() + check(f"{label} max_abs_diff={diff:.2e}", np.allclose(got, want, atol=1e-4)) + +np.random.seed(0) +X = (1, 2, 7, 6) + +# ── conv padding ───────────────────────────────────────────────────────────── + +print("\n=== conv2d padding ===") + +W = numpy_helper.from_array(np.random.randn(3, 2, 3, 3).astype(np.float32), "W") +B = numpy_helper.from_array(np.random.randn(3).astype(np.float32), "B") + +def conv(**attrs): + return helper.make_node("Conv", ["x", "W", "B"], ["y"], kernel_shape=[3, 3], **attrs) + +compare("conv symmetric", conv(pads=[1, 1, 1, 1]), X, [W, B]) +compare("conv asymmetric", conv(pads=[0, 1, 2, 0]), X, [W, B]) +compare("conv SAME_UPPER s2", conv(auto_pad="SAME_UPPER", strides=[2, 2]), X, [W, B]) +compare("conv SAME_LOWER s2", conv(auto_pad="SAME_LOWER", strides=[2, 2]), X, [W, B]) +compare("conv VALID", conv(auto_pad="VALID"), X, [W, B]) + +# ── pooling ────────────────────────────────────────────────────────────────── + +print("\n=== pooling ===") + +def pool(op, **attrs): + return helper.make_node(op, ["x"], ["y"], kernel_shape=[3, 3], **attrs) + +compare("maxpool asymmetric", pool("MaxPool", pads=[0, 1, 2, 0]), X) +compare("maxpool SAME_UPPER s2", pool("MaxPool", auto_pad="SAME_UPPER", strides=[2, 2]), X) +compare("maxpool ceil_mode", pool("MaxPool", strides=[2, 2], ceil_mode=1), X) +compare("avgpool asymmetric incl pad", + pool("AveragePool", pads=[0, 1, 2, 0], count_include_pad=1), X) +compare("avgpool asymmetric excl pad", + pool("AveragePool", pads=[0, 1, 2, 0], count_include_pad=0), X) +compare("avgpool SAME_LOWER s2", + pool("AveragePool", auto_pad="SAME_LOWER", strides=[2, 2]), X) + +# ── reshape ────────────────────────────────────────────────────────────────── + +print("\n=== reshape ===") + +def reshape(shape): + s = numpy_helper.from_array(np.array(shape, dtype=np.int64), "s") + return helper.make_node("Reshape", ["x", "s"], ["y"]), [s] + +node, init = reshape([0, -1]) +compare("reshape [0, -1]", node, (2, 3, 4), init) +node, init = reshape([0, 0, -1]) +compare("reshape [0, 0, -1]", node, (2, 3, 4), init) +node, init = reshape([-1, 4]) +compare("reshape [-1, 4]", node, (2, 3, 4), init) + +# ── load errors ────────────────────────────────────────────────────────────── + +print("\n=== load errors ===") + +path = save_model("unsupported", helper.make_node("Exp", ["x"], ["y"]), (2, 3)) +try: + burn.load_onnx(path) + check("unsupported op raises at load", False) +except RuntimeError as e: + check(f"unsupported op raises at load ({e})", "Exp" in str(e)) + +print() + +if failures: + print(f"{failures} check(s) failed") + sys.exit(1)