Add option to treat initializers as non-constants

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:

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

onnxoptimizer/optimize.h

@@ -94,6 +94,18 @@ const std::vector<std::string> GetAvailablePasses();
 
 const std::vector<std::string> 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<std::string> &names);
 

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<int64_t>(t, val);

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 <typename W, typename... Args>
@@ -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;

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"])