From b12fb2bad2646ddb35de595804a44ba8f2a270be Mon Sep 17 00:00:00 2001 From: Nanmur <1067841003@qq.com> Date: Sat, 12 Sep 2026 20:08:10 +0800 Subject: [PATCH] [ONNX] Fix initializer prefix stripping The ONNX importer used str.strip("onnx::") when parameters were kept as inputs. Since strip treats its argument as a set of characters, it could truncate legitimate initializer names such as those beginning with "n". Use removeprefix so only an exact "onnx::" prefix is removed, and add regression coverage for both prefixed and ordinary names. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 2 +- tests/python/relax/test_frontend_onnx.py | 30 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 62cde1ee12cf..783937be50d9 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -6339,7 +6339,7 @@ def _parse_graph_initializers(self, graph: onnx.onnx_ml_pb2.GraphProto): # Create variables for constants. if self._keep_params_in_input: # Pytorch sometimes inserts silly weight prefix. Remove it. - var_name = init_tensor.name.strip("onnx::") + var_name = init_tensor.name.removeprefix("onnx::") init_var = self._new_var(var_name, shape=array.shape, dtype=array.dtype) self._nodes[init_tensor.name] = init_var # We need to keep track of both the real value and variable for this variable. diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 92f62d32f66d..73fe2d830b9d 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -11767,6 +11767,36 @@ def main( tvm.ir.assert_structural_equal(tvm_model, Expected) +@pytest.mark.parametrize( + ("initializer_name", "expected_name"), + [ + ("onnx::weight", "weight"), + ( + "neck.lateral_convs.2.conv2.weight_quantized", + "neck.lateral_convs.2.conv2.weight_quantized", + ), + ], +) +def test_initializer_name_only_removes_onnx_prefix(initializer_name, expected_name): + graph = helper.make_graph( + [helper.make_node("Add", ["input", initializer_name], ["output"])], + "test_initializer_name_only_removes_onnx_prefix", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], + initializer=[numpy_helper.from_array(np.ones([1], dtype="float32"), initializer_name)], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 14)]) + model.ir_version = 8 + + tvm_model = from_onnx( + model, + keep_params_in_input=True, + sanitize_input_names=False, + ) + + assert tvm_model["main"].params[-1].name == expected_name + + def test_shape_dim_string_expression_graph_add(): identity_node = helper.make_node("Identity", ["x"], ["y"])