Traning mode removal from ONNX nodes

Signed-off-by: Riyad Islam <[email protected]>

Author: i-riyad

modelopt/onnx/utils.py

@@ -25,7 +25,6 @@
 import numpy as np
 import onnx
 import onnx_graphsurgeon as gs
-from onnx import TensorProto, ValueInfoProto, numpy_helper
 from onnx.helper import get_attribute_value
 from onnx_graphsurgeon import Constant, Node, Variable
 
@@ -289,7 +288,7 @@ def _convert_types_to_np(types: dict[str, int] | list[int] | int) -> Any:
 
 def get_tensor_by_name(
     onnx_model: onnx.ModelProto, tensor_name: str
-) -> ValueInfoProto | TensorProto | None:
+) -> onnx.ValueInfoProto | onnx.TensorProto | None:
     """This function returns a tensor from its name.
 
     This function searches for a tensor in the model's:
@@ -438,7 +437,7 @@ def randomize_weights_onnx_bytes(onnx_bytes: bytes, seed: int = 0) -> bytes:
                     numpy_array = np.random.normal(float(avg), float(var), size=init.dims).astype(
                         dtype
                     )
-                    tensor = numpy_helper.from_array(numpy_array, init.name)
+                    tensor = onnx.numpy_helper.from_array(numpy_array, init.name)
                     model.graph.initializer[idx].CopyFrom(tensor)
 
     buffer = io.BytesIO()
@@ -751,3 +750,46 @@ def onnx_type_str_to_enum(dtype: str) -> int:
     dtype = dtype.split("tensor(")[-1].split(")")[0]
     dtype = "FLOAT" if dtype == "float32" else dtype.upper()
     return getattr(onnx.TensorProto, dtype)
+
+
+def remove_node_training_mode(onnx_model: onnx.ModelProto, node_op_type: str) -> onnx.ModelProto:
+    """Remove `training_mode` attribute and extra training outputs from nodes of a given op type.
+
+    Args:
+        onnx_model: The onnx model.
+        node_op_type: The node type to remove training_mode attribute from.
+
+    Returns:
+        The onnx model with the training_mode attribute removed.
+    """
+    removed_output_names = set()
+
+    for node in onnx_model.graph.node:
+        if node.op_type != node_op_type:
+            continue
+
+        # Drop the 'training_mode' attribute if present
+        for idx, attr in enumerate(list(node.attribute)):
+            if attr.name == "training_mode":
+                del node.attribute[idx]
+                break
+
+        # If node has extra training outputs, keep only the first
+        if len(node.output) > 1:
+            removed_output_names.update(node.output[1:])
+            node.output[:] = node.output[:1]
+
+    if removed_output_names:
+        # Clean up corresponding value_info entries
+        keep = [vi for vi in onnx_model.graph.value_info if vi.name not in removed_output_names]
+        del onnx_model.graph.value_info[:]
+        onnx_model.graph.value_info.extend(keep)
+
+        # Also clean up graph.output entries
+        keep_outputs = [
+            out for out in onnx_model.graph.output if out.name not in removed_output_names
+        ]
+        del onnx_model.graph.output[:]
+        onnx_model.graph.output.extend(keep_outputs)
+
+    return onnx_model

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -45,6 +45,7 @@
     get_node_names,
     get_output_names,
     get_output_shapes,
+    remove_node_training_mode,
 )
 from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers
 from modelopt.torch.utils import flatten_tree, standardize_named_model_args
@@ -569,25 +570,3 @@ def get_onnx_bytes(*args, **kwargs) -> bytes:
     onnx_bytes = get_onnx_bytes_and_metadata(*args, **kwargs)[0]
     onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes)
     return onnx_bytes_obj.get_onnx_model_file_bytes()
-
-
-def remove_node_training_mode(onnx_model: ModelProto, node_op_type: str) -> ModelProto:
-    """Remove training_mode attribute from selected node type.
-
-    Args:
-        onnx_model: The onnx model.
-        node_op_type: The node type to remove training_mode attribute from.
-
-    Returns:
-        The onnx model with the training_mode attribute removed.
-    """
-    for node in onnx_model.graph.node:
-        if node.op_type == node_op_type:
-            for attribute in node.attribute:
-                if attribute.name == "training_mode":
-                    if attribute.i == 1:
-                        node.output.remove(node.output[1])
-                        node.output.remove(node.output[1])
-                    attribute.i = 0
-
-    return onnx_model

tests/unit/torch/deploy/utils/test_torch_onnx_utils.py

@@ -24,13 +24,21 @@
 import torch.nn as nn
 from _test_utils.torch_model.deploy_models import BaseDeployModel, get_deploy_models
 from _test_utils.torch_model.vision_models import get_tiny_resnet_and_input
-from onnx.helper import make_graph, make_model, make_node, make_tensor_value_info
+from onnx.helper import (
+    make_graph,
+    make_model,
+    make_node,
+    make_opsetid,
+    make_tensor,
+    make_tensor_value_info,
+)
 
 from modelopt.onnx.utils import (
     get_batch_size_from_bytes,
     get_input_names_from_bytes,
     get_output_names_from_bytes,
     randomize_weights_onnx_bytes,
+    remove_node_training_mode,
     remove_weights_data,
     validate_batch_size,
 )
@@ -255,3 +263,80 @@ def test_reproducible_random_weights():
     onnx_bytes_1 = randomize_weights_onnx_bytes(onnx_bytes_wo_weights)
     onnx_bytes_2 = randomize_weights_onnx_bytes(onnx_bytes_wo_weights)
     assert onnx_bytes_1 == onnx_bytes_2
+
+
+def _make_bn_initializer(name: str, shape, value=1.0):
+    """Helper to create an initializer tensor for BatchNorm."""
+    data = np.full(shape, value, dtype=np.float32)
+    return make_tensor(name, onnx.TensorProto.FLOAT, shape, data.flatten())
+
+
+def _make_batchnorm_model(bn_node, extra_value_infos=None):
+    """Helper to create an ONNX model with a BatchNormalization node."""
+    initializers = [
+        _make_bn_initializer("scale", [3], 1.0),
+        _make_bn_initializer("bias", [3], 0.0),
+        _make_bn_initializer("mean", [3], 0.0),
+        _make_bn_initializer("var", [3], 1.0),
+    ]
+
+    graph_def = make_graph(
+        [bn_node],
+        "test_graph",
+        [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1, 3, 224, 224])],
+        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1, 3, 224, 224])],
+        initializer=initializers,
+        value_info=extra_value_infos or [],
+    )
+
+    return make_model(graph_def, opset_imports=[make_opsetid("", 14)])
+
+
+def test_remove_node_training_mode_attribute():
+    """Test removal of training_mode attribute from BatchNormalization nodes."""
+    bn_node = make_node(
+        "BatchNormalization",
+        inputs=["input", "scale", "bias", "mean", "var"],
+        outputs=["output"],
+        name="bn1",
+        training_mode=1,  # This attribute should be removed
+    )
+
+    model = _make_batchnorm_model(bn_node)
+    result_model = remove_node_training_mode(model, "BatchNormalization")
+
+    bn_node_result = result_model.graph.node[0]
+    assert bn_node_result.op_type == "BatchNormalization"
+
+    # Check that training_mode attribute is not present
+    attr_names = [attr.name for attr in bn_node_result.attribute]
+    assert "training_mode" not in attr_names
+
+
+def test_remove_node_extra_training_outputs():
+    """Test removal of extra training outputs from BatchNormalization nodes."""
+    bn_node = make_node(
+        "BatchNormalization",
+        inputs=["input", "scale", "bias", "mean", "var"],
+        outputs=["output", "saved_mean", "saved_inv_std"],  # Extra training outputs
+        name="bn1",
+        training_mode=1,
+    )
+
+    value_infos = [
+        make_tensor_value_info("saved_mean", onnx.TensorProto.FLOAT, [3]),
+        make_tensor_value_info("saved_inv_std", onnx.TensorProto.FLOAT, [3]),
+    ]
+
+    model = _make_batchnorm_model(bn_node, extra_value_infos=value_infos)
+    result_model = remove_node_training_mode(model, "BatchNormalization")
+
+    # Verify only first output remains
+    bn_node_result = result_model.graph.node[0]
+    assert len(bn_node_result.output) == 1
+    assert bn_node_result.output[0] == "output"
+
+    # Verify value_info entries for removed outputs are cleaned up
+    value_info_names = [vi.name for vi in result_model.graph.value_info]
+    assert "saved_mean" not in value_info_names
+    assert "saved_inv_std" not in value_info_names