This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d3a02fbef6 [ONNX] Fix initializer prefix stripping (#20323)
d3a02fbef6 is described below
commit d3a02fbef68a3cc744458e4fb5eab0b2e3730e14
Author: Nanmur <[email protected]>
AuthorDate: Mon Sep 14 12:05:19 2026 +0800
[ONNX] Fix initializer prefix stripping (#20323)
The Relax ONNX importer used `str.strip("onnx::")` when
`keep_params_in_input=True`. Since `strip` treats its argument as a set
of characters, initializer names that do not have the prefix can still
lose leading or trailing `o`, `n`, `x`, or `:` characters. For example,
`neck...` became `eck...`.
This change uses `str.removeprefix("onnx::")` so only the exact
PyTorch-generated prefix is removed. The regression test covers both a
prefixed initializer and an ordinary initializer beginning with `n`.
Fixes #20290
Tests:
- `python -m pytest tests/python/relax/test_frontend_onnx.py -k
'params_names_start_with_onnx or
initializer_name_only_removes_onnx_prefix or concat_with_param or
multi_ops_with_same_params' -q`
- `python -m ruff check python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
- `python -m ruff format --check
python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
Local full-file note: the complete ONNX test file reached 28 passing
tests before hitting the unrelated Windows JIT export assertion `Target
triple should not be empty`.
---
python/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 498d0070fe..93a1734166 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -6354,7 +6354,7 @@ class ONNXGraphImporter:
# 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 660179f3b4..2e07dd7908 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -11797,6 +11797,36 @@ def test_params_names_start_with_onnx():
tvm.ir.assert_structural_equal(tvm_model, Expected)
[email protected](
+ ("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"])
