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 7356265096 [Fix][Relax][ONNX] Cast BatchNorm params to input dtype
(#19979)
7356265096 is described below
commit 7356265096cba7196673b09f90b023e600171452
Author: Vic Wen <[email protected]>
AuthorDate: Sat Jul 11 14:54:13 2026 +0800
[Fix][Relax][ONNX] Cast BatchNorm params to input dtype (#19979)
Fixes #19977.
ONNX `BatchNormalization` allows the input/output tensor dtype,
scale/bias dtype, and mean/variance dtype to be separate floating-point
type parameters.
For example, a valid ONNX model may use `float16` data with `float32`
gamma, beta, mean, and variance tensors.
The Relax `batch_norm` operator currently requires all five input
tensors to have the same dtype. The ONNX frontend previously forwarded
the ONNX inputs directly to `relax.nn.batch_norm`, causing import to
fail during normalization
for mixed-dtype ONNX models.
This patch casts the ONNX BatchNormalization parameter tensors (`scale`,
`bias`, `mean`, and `var`) to the data tensor dtype before calling Relax
`batch_norm`.
This preserves the ONNX output dtype, which follows the input data
dtype, while keeping the fix localized to the frontend compatibility
layer.
The regression test builds a minimal ONNX BatchNormalization graph with
`float16` data and `float32` parameters, imports it through the Relax
ONNX frontend, and checks that the generated Relax `batch_norm` call
receives same-dtype inputs.
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_batch_norm_mixed_dtype_params
tests/python/relax/test_frontend_onnx.py::test_batch_norm_defaults_to_inference_mode
-q`
Signed-off-by: viiccwen <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 69 ++++++++++++++++--
tests/python/relax/test_frontend_onnx.py | 95 +++++++++++++++++++++++++
2 files changed, 158 insertions(+), 6 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index be3dd5d4e4..da65cc7bee 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -3576,18 +3576,75 @@ class BatchNormalization(OnnxOpConverter):
epsilon = attr.get("epsilon", 1e-05)
momentum = attr.get("momentum", 0.9)
training_mode = attr.get("training_mode", 0)
- return relax.op.nn.batch_norm(
- data,
- gamma=scale,
- beta=bias,
- moving_mean=mean,
- moving_var=var,
+
+ data_dtype = data.ty.dtype
+ scale_dtype = scale.ty.dtype
+ bias_dtype = bias.ty.dtype
+ mean_dtype = mean.ty.dtype
+ var_dtype = var.ty.dtype
+
+ if scale_dtype != bias_dtype:
+ raise ValueError(
+ "ONNX BatchNormalization requires scale and bias to have the
same "
+ f"dtype, but received {scale_dtype} and {bias_dtype}."
+ )
+
+ if mean_dtype != var_dtype:
+ raise ValueError(
+ "ONNX BatchNormalization requires mean and var to have the
same "
+ f"dtype, but received {mean_dtype} and {var_dtype}."
+ )
+
+ if data_dtype == scale_dtype == mean_dtype:
+ compute_dtype = data_dtype
+ elif (
+ data_dtype == "float16"
+ and scale_dtype in ("float16", "float32")
+ and mean_dtype in ("float16", "float32")
+ ):
+ compute_dtype = "float32"
+ else:
+ raise NotImplementedError(
+ "ONNX BatchNormalization with mixed input dtypes is currently "
+ "supported only for float16 data with float16/float32
parameters "
+ "and statistics, but received "
+ f"data={data_dtype}, scale/bias={scale_dtype},
mean/var={mean_dtype}."
+ )
+
+ # ONNX requires float computation for float16 training statistics to
avoid overflow.
+ if training_mode and data_dtype == "float16":
+ compute_dtype = "float32"
+
+ def cast_for_compute(expr, source_dtype):
+ if source_dtype == compute_dtype:
+ return expr
+ return relax.op.astype(expr, compute_dtype)
+
+ output = relax.op.nn.batch_norm(
+ cast_for_compute(data, data_dtype),
+ gamma=cast_for_compute(scale, scale_dtype),
+ beta=cast_for_compute(bias, bias_dtype),
+ moving_mean=cast_for_compute(mean, mean_dtype),
+ moving_var=cast_for_compute(var, var_dtype),
axis=1,
epsilon=epsilon,
momentum=momentum,
training=bool(training_mode),
)
+ y = relax.TupleGetItem(output, 0)
+ running_mean = relax.TupleGetItem(output, 1)
+ running_var = relax.TupleGetItem(output, 2)
+
+ if compute_dtype != data_dtype:
+ y = relax.op.astype(y, data_dtype)
+ if compute_dtype != mean_dtype:
+ running_mean = relax.op.astype(running_mean, mean_dtype)
+ if compute_dtype != var_dtype:
+ running_var = relax.op.astype(running_var, var_dtype)
+
+ return relax.Tuple([y, running_mean, running_var])
+
class MeanVarianceNormalization(OnnxOpConverter):
"""Converts an onnx MeanVarianceNormalization node into an equivalent
Relax expression."""
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 5aab4f558f..35cf699dba 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -7546,6 +7546,101 @@ def test_batch_norm_defaults_to_inference_mode():
assert batch_norm_attrs[0].training is False
+def test_batch_norm_mixed_dtype_params():
+ data = helper.make_tensor_value_info("data", TensorProto.FLOAT16, [1, 3,
2, 2])
+ output = helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1,
3, 2, 2])
+ params = [
+ numpy_helper.from_array(np.array([1.0, 1.5, 2.0], dtype=np.float32),
name="gamma"),
+ numpy_helper.from_array(np.array([0.0, 0.1, -0.1], dtype=np.float32),
name="beta"),
+ numpy_helper.from_array(np.array([0.2, -0.3, 0.4], dtype=np.float32),
name="mean"),
+ numpy_helper.from_array(np.array([1.0, 1.5, 2.0], dtype=np.float32),
name="var"),
+ ]
+ batch_norm_node = helper.make_node(
+ "BatchNormalization",
+ ["data", "gamma", "beta", "mean", "var"],
+ ["output"],
+ epsilon=1e-5,
+ momentum=0.9,
+ training_mode=0,
+ )
+ graph = helper.make_graph(
+ [batch_norm_node],
+ "mixed_dtype_batchnorm",
+ [data],
+ [output],
+ initializer=params,
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
15)])
+
+ tvm_model = from_onnx(model, keep_params_in_input=False)
+
+ assert tuple(dim.value for dim in tvm_model["main"].ret_ty.shape.values)
== (1, 3, 2, 2)
+ assert tvm_model["main"].ret_ty.dtype == "float16"
+
+ batch_norm_calls = []
+
+ def visit(expr):
+ if isinstance(expr, relax.Call) and expr.op ==
tvm.ir.Op.get("relax.nn.batch_norm"):
+ batch_norm_calls.append(expr)
+
+ relax.analysis.post_order_visit(tvm_model["main"], visit)
+
+ assert len(batch_norm_calls) == 1
+ arg_dtypes = [
+ str(getattr(arg, "struct_info", getattr(arg, "ty", None)).dtype)
+ for arg in batch_norm_calls[0].args
+ ]
+ assert arg_dtypes == ["float32"] * 5
+
+
+def test_batch_norm_training_preserves_output_dtypes():
+ data = helper.make_tensor_value_info("data", TensorProto.FLOAT16, [1, 3,
2, 2])
+ outputs = [
+ helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1, 3, 2,
2]),
+ helper.make_tensor_value_info("running_mean", TensorProto.FLOAT16,
[3]),
+ helper.make_tensor_value_info("running_var", TensorProto.FLOAT16, [3]),
+ ]
+ inputs = [
+ data,
+ helper.make_tensor_value_info("gamma", TensorProto.FLOAT16, [3]),
+ helper.make_tensor_value_info("beta", TensorProto.FLOAT16, [3]),
+ helper.make_tensor_value_info("mean", TensorProto.FLOAT16, [3]),
+ helper.make_tensor_value_info("var", TensorProto.FLOAT16, [3]),
+ ]
+ batch_norm_node = helper.make_node(
+ "BatchNormalization",
+ [value.name for value in inputs],
+ [value.name for value in outputs],
+ training_mode=1,
+ )
+ graph = helper.make_graph(
+ [batch_norm_node],
+ "mixed_dtype_training_batchnorm",
+ inputs,
+ outputs,
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
15)])
+
+ tvm_model = from_onnx(model, keep_params_in_input=True)
+
+ assert [str(field.dtype) for field in tvm_model["main"].ret_ty.fields] == [
+ "float16",
+ "float16",
+ "float16",
+ ]
+
+ batch_norm_calls = []
+
+ def visit(expr):
+ if isinstance(expr, relax.Call) and expr.op ==
tvm.ir.Op.get("relax.nn.batch_norm"):
+ batch_norm_calls.append(expr)
+
+ relax.analysis.post_order_visit(tvm_model["main"], visit)
+
+ assert len(batch_norm_calls) == 1
+ assert [str(arg.ty.dtype) for arg in batch_norm_calls[0].args] ==
["float32"] * 5
+
+
def get_pool_padding(shape, auto_pad, kernel_shape, strides, pads):
def get_pad_pair(input1d, kernel1d, stride1d, mode):
if input1d % stride1d == 0:
