From c8855bd7580623580db6819451844db2a1c22ee4 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 28 Jul 2026 06:45:51 +0000 Subject: [PATCH] Add option to treat initializers as non-constants Introduce a thread-local switch, SetInitializersAsConstants / InitializersAsConstants, that controls whether the fusion/elimination passes treat graph initializers as constant tensors. The single choke points IsConstantTensor and FetchConstantTensor consult it, so when the switch is off every value-baking pass (fuse_bn_into_conv, nop-reshape on a constant shape, ...) leaves initializer-backed weights untouched while Constant nodes stay constant. Expose it through the Python binding and as an initializers_as_constants keyword on onnxoptimizer.optimize(), which sets the switch around the run and restores it afterwards. Add tests covering the fuse-bn case and that the switch is restored. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MyzBiNk5LvTaVPgURqD5UH --- onnxoptimizer/__init__.py | 26 ++++++++++- onnxoptimizer/cpp2py_export.cc | 6 +++ onnxoptimizer/optimize.h | 12 +++++ onnxoptimizer/passes/pass_util.cc | 16 +++++++ onnxoptimizer/passes/pass_util.h | 17 ++++++- onnxoptimizer/test/optimizer_test.py | 68 ++++++++++++++++++++++++++++ 6 files changed, 142 insertions(+), 3 deletions(-) diff --git a/onnxoptimizer/__init__.py b/onnxoptimizer/__init__.py index 5af997aa5..81726d27e 100644 --- a/onnxoptimizer/__init__.py +++ b/onnxoptimizer/__init__.py @@ -33,13 +33,23 @@ def optimize( - model: onnx.ModelProto, passes: list[str] | None = None, fixed_point: bool = False + model: onnx.ModelProto, + passes: list[str] | None = None, + fixed_point: bool = False, + initializers_as_constants: bool = True, ) -> onnx.ModelProto: """Apply the optimization on the serialized ModelProto. Arguments: model: ONNX model. passes: Optimization names. + fixed_point: Whether to run the passes to a fixed point. + initializers_as_constants: Whether the passes may treat graph + initializers as constant tensors (the default, ``True``). When set + to ``False`` initializers are treated as non-constant, so + value-baking passes such as ``fuse_bn_into_conv`` leave + initializer-backed weights untouched; ``Constant`` nodes are still + treated as constants. Return: Optimized model. @@ -49,6 +59,20 @@ def optimize( passes = get_fuse_and_elimination_passes() if not isinstance(model, onnx.ModelProto): raise TypeError(f"Optimizer only accepts ModelProto, incorrect type: {type(model)}") + # The C++ core reads this switch from thread-local state deep inside the + # passes, so set it around the call and restore it afterwards to avoid + # leaking the setting to unrelated callers on the same thread. + previous = _c.initializers_as_constants() + _c.set_initializers_as_constants(initializers_as_constants) + try: + return _optimize_impl(model, passes, fixed_point) + finally: + _c.set_initializers_as_constants(previous) + + +def _optimize_impl( + model: onnx.ModelProto, passes: list[str], fixed_point: bool +) -> onnx.ModelProto: try: model_str = model.SerializeToString() if fixed_point: diff --git a/onnxoptimizer/cpp2py_export.cc b/onnxoptimizer/cpp2py_export.cc index 5b9c27a29..55169575c 100644 --- a/onnxoptimizer/cpp2py_export.cc +++ b/onnxoptimizer/cpp2py_export.cc @@ -76,5 +76,11 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) { &optimization::GetAvailablePasses); onnx_opt_cpp2py_export.def("get_fuse_and_elimination_passes", &optimization::GetFuseAndEliminationPass); + // Toggle whether the passes treat graph initializers as constant tensors + // (default true). See SetInitializersAsConstants in optimize.h. + onnx_opt_cpp2py_export.def("set_initializers_as_constants", + &optimization::SetInitializersAsConstants); + onnx_opt_cpp2py_export.def("initializers_as_constants", + &optimization::InitializersAsConstants); } } // namespace ONNX_NAMESPACE diff --git a/onnxoptimizer/optimize.h b/onnxoptimizer/optimize.h index 6d59429e0..fd1983f30 100644 --- a/onnxoptimizer/optimize.h +++ b/onnxoptimizer/optimize.h @@ -94,6 +94,18 @@ const std::vector GetAvailablePasses(); const std::vector GetFuseAndEliminationPass(); +// Control whether the optimizer passes treat graph initializers as constant +// tensors. The default (true) is onnxoptimizer's historical behaviour, in which +// an initializer-backed value is a constant and value-baking passes +// (fuse_bn_into_conv, fuse_add_bias_into_conv, nop-reshape/expand on a constant +// shape, ...) may consume and fold it. When set to false, initializers are +// treated as non-constant, so those passes leave initializer-backed values -- +// and the weights they represent -- untouched; Constant *nodes* are still +// treated as constants. The setting is thread-local and stays in effect until +// changed, so callers that flip it should restore it afterwards. +void SetInitializersAsConstants(bool value); +bool InitializersAsConstants(); + ModelProto Optimize(const ModelProto &mp_in, const std::vector &names); diff --git a/onnxoptimizer/passes/pass_util.cc b/onnxoptimizer/passes/pass_util.cc index e0c3d9c37..f99463651 100644 --- a/onnxoptimizer/passes/pass_util.cc +++ b/onnxoptimizer/passes/pass_util.cc @@ -11,6 +11,22 @@ namespace ONNX_NAMESPACE { namespace optimization { +// Process-wide (per-thread) switch controlling whether the fusion/elimination +// passes treat graph initializers as constant tensors. Defaults to true, which +// is onnxoptimizer's historical behaviour. See the declarations in +// ``optimize.h`` for the full contract. +namespace { +thread_local bool g_initializers_as_constants = true; +} // namespace + +bool InitializersAsConstants() { + return g_initializers_as_constants; +} + +void SetInitializersAsConstants(bool value) { + g_initializers_as_constants = value; +} + bool FetchSoleIntValueOfTensor(const Value* t, int64_t& val) { int32_t i32_val; const bool r1 = FetchSoleValueOfTensor(t, val); diff --git a/onnxoptimizer/passes/pass_util.h b/onnxoptimizer/passes/pass_util.h index 759127084..143a94f30 100644 --- a/onnxoptimizer/passes/pass_util.h +++ b/onnxoptimizer/passes/pass_util.h @@ -119,9 +119,22 @@ T1 AddYIfNegative(T1 x, T2 y) { return x < 0 ? x + y : x; } +// Whether the fusion/elimination passes currently treat graph initializers as +// constant tensors. Toggled with SetInitializersAsConstants (declared in +// optimize.h); declared here because the inline constant helpers below consult +// it. Defaults to true (historical behaviour). +bool InitializersAsConstants(); + inline bool IsConstantTensor(const Value* v) { auto* graph = v->owningGraph(); - return v->node()->kind() == kConstant || graph->is_constant_initializer(v); + if (v->node()->kind() == kConstant) { + return true; + } + // When initializers are treated as non-constant, a value backed only by an + // initializer is not a constant, so value-baking passes (fuse_bn_into_conv, + // nop-reshape on a constant shape, ...) leave it -- and the weight it + // represents -- untouched. Constant *nodes* stay constant either way. + return InitializersAsConstants() && graph->is_constant_initializer(v); } template @@ -140,7 +153,7 @@ inline const Tensor* FetchConstantTensor(const Value* v) { auto* graph = v->owningGraph(); if (kind == kConstant && v->node()->hasAttribute(kvalue)) { return &v->node()->t(kvalue); - } else if (graph->is_constant_initializer(v)) { + } else if (InitializersAsConstants() && graph->is_constant_initializer(v)) { return &*graph->getInitializer(v->uniqueName()); } else { return nullptr; diff --git a/onnxoptimizer/test/optimizer_test.py b/onnxoptimizer/test/optimizer_test.py index 2ee7d3003..b65436e88 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/onnxoptimizer/test/optimizer_test.py @@ -3282,6 +3282,74 @@ def test_fuse_bn_into_conv_simple(self): # type: () -> None ) optimized_model = self._optimized(graph, ["fuse_bn_into_conv"]) # noqa + def _make_conv_bn_model(self): # type: () -> ModelProto + tensor_type, np_type = TensorProto.FLOAT, np.float32 + conv = helper.make_node("Conv", ["X", "W", "B"], ["Y"]) + bn = helper.make_node( + "BatchNormalization", ["Y", "scale", "b", "mean", "var"], ["Z"] + ) + W = np.random.randn(3, 2, 5, 5).astype(np_type) + 2 + B = np.random.randn(3).astype(np_type) + 2 + scale = np.random.randn(3).astype(np_type) + 2 + b = np.random.randn(3).astype(np_type) + 2 + mean = np.random.randn(3).astype(np_type) + 2 + var = np.abs(np.random.randn(3).astype(np_type)) + 2 + initializers = [ + helper.make_tensor(name, tensor_type, npa.shape, npa.tobytes(), raw=True) + for name, npa in [ + ("W", W), + ("B", B), + ("scale", scale), + ("b", b), + ("mean", mean), + ("var", var), + ] + ] + graph = helper.make_graph( + [conv, bn], + "test", + [helper.make_tensor_value_info("X", tensor_type, (5, 2, 28, 28))], + [helper.make_tensor_value_info("Z", tensor_type, (5, 3, 24, 24))], + initializer=initializers, + value_info=[ + helper.make_tensor_value_info("Y", tensor_type, (5, 3, 24, 24)) + ], + ) + return helper.make_model( + graph, + producer_name="onnx-test", + opset_imports=[helper.make_opsetid("", LATEST_STABLE_OPSET_VERSION)], + ir_version=10, + ) + + def test_fuse_bn_into_conv_default_treats_initializers_as_constants(self): + # With the default behaviour the BatchNormalization initializers are + # constants, so fuse_bn_into_conv folds the BN into the Conv weights. + model = self._make_conv_bn_model() + optimized = onnxoptimizer.optimize(model, ["fuse_bn_into_conv"]) + op_types = [n.op_type for n in optimized.graph.node] + assert "BatchNormalization" not in op_types + + def test_initializers_as_non_constants_disables_fuse_bn(self): + # Treating initializers as non-constant leaves the BN weights alone, so + # the pass cannot fold BatchNormalization into the Conv. + model = self._make_conv_bn_model() + optimized = onnxoptimizer.optimize( + model, ["fuse_bn_into_conv"], initializers_as_constants=False + ) + op_types = [n.op_type for n in optimized.graph.node] + assert "BatchNormalization" in op_types + assert "Conv" in op_types + + def test_initializers_as_constants_flag_is_restored(self): + # onnxoptimizer.optimize must not leak the thread-local switch it sets. + assert onnxoptimizer.onnx_opt_cpp2py_export.initializers_as_constants() + model = self._make_conv_bn_model() + onnxoptimizer.optimize( + model, ["fuse_bn_into_conv"], initializers_as_constants=False + ) + assert onnxoptimizer.onnx_opt_cpp2py_export.initializers_as_constants() + def _internal_test_deadend_elimination(self, fixed): # type: (bool) -> None softmax = helper.make_node("Softmax", ["X"], ["Y"], axis=2) log = helper.make_node("Log", ["Y"], ["Z"])