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 9fbdd94c3b [Fix][Relax][Frontend][TFLite] Correct quantized SSD
inference (#20291)
9fbdd94c3b is described below
commit 9fbdd94c3b67a81a7edc88182aa4d28412e8b4e2
Author: Hongyi Wu <[email protected]>
AuthorDate: Fri Sep 11 11:43:16 2026 +0800
[Fix][Relax][Frontend][TFLite] Correct quantized SSD inference (#20291)
## Summary
This PR fixes Relax import and execution of quantized TFLite SSD models
and adds focused regression coverage for the repaired conversion paths.
- Preserve TFLite `DETECTION_POSTPROCESS` class semantics by removing
the optional background channel, keeping zero-based labels, and avoiding
a second softmax.
- Add quantized `AVERAGE_POOL_2D` and `MAX_POOL_2D` to the supported QDQ
conversion path.
- Correct per-channel depthwise-convolution dequantization and the
weight layout when `depth_multiplier > 1`.
- Add targeted frontend, operator, and parser regressions for those
conversion semantics.
## Design
### TFLite detection postprocessing
TFLite SSD class scores are already activated. When the input class
tensor has one more channel than the custom op's `num_classes`, that
extra channel is the background class. The frontend now validates that
the difference is either zero or one, slices the background channel when
present, and keeps the remaining class indices zero-based.
`relax.vision.multibox_transform_loc` gains an `apply_softmax`
attribute. Its default remains `True`, preserving the existing generic
Relax behavior, while the TFLite frontend sets it to `False` for
`DETECTION_POSTPROCESS`.
### Quantized pooling and depthwise convolution
Average and max pooling now use the same QDQ conversion path as the
other supported quantized operators.
For per-channel depthwise weights, dequantization is applied along the
original TFLite channel axis before the weights are reshaped to Relax's
convolution layout. The reshape also distinguishes the `input_channels
== 1` case from the general `depth_multiplier > 1` case, producing the
correct HWIO or HWOI layout respectively.
## Review map
1. `python/tvm/relax/frontend/tflite/tflite_frontend.py`: TFLite
detection, pooling, and depthwise conversion.
2. `include/tvm/relax/attrs/vision.h`, `src/relax/op/vision/`, and the
Python Relax/TOPI vision layers: `apply_softmax` plumbing with the
backward-compatible default.
3. `tests/python/relax/test_frontend_tflite.py`: focused
detection-postprocess, quantized-pool, and per-channel-depthwise
frontend regressions.
4. `tests/python/relax/test_op_vision.py` and
`tests/python/relax/test_tvmscript_parser_op_vision.py`: operator and
parser coverage.
## Validation
Validated locally after the final amended commit.
| Environment | Result |
| --- | --- |
| Local external-disk build | Full build passed |
| Local TFLite frontend suite | 559 passed |
| Local Relax vision operator and parser suites | 70 passed |
| Local formatting/static checks | Ruff and clang-format checks passed |
The official quantized SSD MobileNet V2 COCO model checkpoint/artifact
was also imported, compiled, and executed locally as an out-of-tree
development validation; it is intentionally not part of the TVM test
suite.
## Known limits
This PR is correctness-focused and retains the existing QDQ lowering
(`int -> float -> int`); it does not claim a performance improvement. In
the full SSD network, accumulated differences between TFLite fixed-point
kernels and Relax QDQ rounding can reorder low-confidence tail
detections. The local model check verifies the leading semantic
detection; integer-native quantized lowering is left as follow-up work.
---
include/tvm/relax/attrs/vision.h | 5 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 100 +++++++++++++++++----
.../tvm/relax/op/vision/multibox_transform_loc.py | 8 +-
python/tvm/relax/transform/legalize_ops/vision.py | 1 +
python/tvm/topi/vision/multibox_transform_loc.py | 13 +--
src/relax/op/vision/multibox_transform_loc.cc | 8 +-
src/relax/op/vision/multibox_transform_loc.h | 2 +-
tests/python/relax/test_frontend_tflite.py | 97 +++++++++++++++++---
tests/python/relax/test_op_vision.py | 17 +++-
.../relax/test_tvmscript_parser_op_vision.py | 2 +
10 files changed, 208 insertions(+), 45 deletions(-)
diff --git a/include/tvm/relax/attrs/vision.h b/include/tvm/relax/attrs/vision.h
index f4b1830669..9a2384a494 100644
--- a/include/tvm/relax/attrs/vision.h
+++ b/include/tvm/relax/attrs/vision.h
@@ -158,6 +158,7 @@ struct MultiboxTransformLocAttrs : public AttrsNode {
double threshold;
ffi::Array<double> variances;
bool keep_background;
+ bool apply_softmax;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -170,7 +171,9 @@ struct MultiboxTransformLocAttrs : public AttrsNode {
"(x,y,w,h) scales = TFLite
1/x_scale,1/y_scale,1/w_scale,1/h_scale on "
"encodings. Very large w/h scales can overflow exp in decode.")
.def_ro("keep_background", &MultiboxTransformLocAttrs::keep_background,
- "If false, force output scores[:,0,:] to 0 (background
class).");
+ "If false, force output scores[:,0,:] to 0 (background
class).")
+ .def_ro("apply_softmax", &MultiboxTransformLocAttrs::apply_softmax,
+ "Whether to apply softmax to class predictions before
thresholding.");
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs",
MultiboxTransformLocAttrs, AttrsNode);
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 13d78c4c30..935246cc18 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -100,6 +100,7 @@ class OperatorConverter:
{
"ABS",
"ADD",
+ "AVERAGE_POOL_2D",
"ATAN2",
"CEIL",
"CONCATENATION",
@@ -125,6 +126,7 @@ class OperatorConverter:
"LOGISTIC",
"LOG_SOFTMAX",
"MAXIMUM",
+ "MAX_POOL_2D",
"MEAN",
"MINIMUM",
"MUL",
@@ -4650,6 +4652,13 @@ class OperatorConverter:
# 1 KH KW C(input_c * depth_multiplier)
_, kernel_h, kernel_w, in_channels =
to_int_list(self.get_tensor_shape(weight_tensor))
assert in_channels == input_c * depth_multiplier
+ # Relax grouped convolution expects one input channel per group in
+ # HWOI when input_c > 1. The single-channel case remains HWIO.
+ depthwise_weight_shape = (
+ (kernel_h, kernel_w, input_c, depth_multiplier)
+ if input_c == 1
+ else (kernel_h, kernel_w, in_channels, 1)
+ )
else:
output_channels, kernel_h, kernel_w, in_channels = to_int_list(
self.get_tensor_shape(weight_tensor)
@@ -4689,6 +4698,13 @@ class OperatorConverter:
)
weight_tensor_type_str = self.get_tensor_type_str(weight_tensor_type)
+ per_channel_depthwise = (
+ is_depthwise_conv
+ and input_tensor.qnn_params
+ and weight_tensor.qnn_params
+ and weight_tensor.tensor.Quantization().ScaleLength() > 1
+ )
+
in_expr = self.get_expr(input_tensor_idx)
# TFLite converts float32 models to float16 models by introducing
@@ -4698,9 +4714,8 @@ class OperatorConverter:
if self.has_expr(weight_tensor.tensor_idx):
weight_expr = self.get_expr(weight_tensor.tensor_idx)
if is_depthwise_conv:
- weight_expr = relax.op.reshape(
- weight_expr, (kernel_h, kernel_w, input_c,
depth_multiplier)
- )
+ if not per_channel_depthwise:
+ weight_expr = relax.op.reshape(weight_expr,
depthwise_weight_shape)
else:
weight_expr = relax.op.permute_dims(weight_expr, axes=(1, 2,
3, 0))
else:
@@ -4713,10 +4728,11 @@ class OperatorConverter:
# convolution:
# OC KH KW IC, we require KH KW IC OC (HWIO)
# depthwise convolution:
- # 1 KH KW C(input_c * depth_multiplier), we require
- # KH KW IC M (depth_multiplier) (HWOI)
+ # 1 KH KW C(input_c * depth_multiplier), we require HWIO
+ # [KH,KW,1,M] for one input channel, otherwise HWOI [KH,KW,C*M,1].
if is_depthwise_conv:
- weight_value = weight_value.reshape(kernel_h, kernel_w,
input_c, depth_multiplier)
+ if not per_channel_depthwise:
+ weight_value = weight_value.reshape(depthwise_weight_shape)
else:
weight_value = weight_value.transpose((1, 2, 3, 0))
@@ -4748,11 +4764,39 @@ class OperatorConverter:
# QuantizedDimension() == 0 (OC in original) → axis 3 in HWIO.
weight_axis = weight_tensor.qnn_params["axis"]
if is_depthwise_conv:
- if weight_axis != 0:
- raise tvm.error.OpNotImplemented(
- "Per-channel quantized depthwise convolution is not
supported "
- "because the channel axis changes semantics after the "
- "[1,KH,KW,C*M] → [KH,KW,C,M] reshape."
+ if per_channel_depthwise:
+ if weight_axis != 3:
+ raise tvm.error.OpAttributeInvalid(
+ "Per-channel DepthwiseConv2D weight
QuantizedDimension() must be 3 "
+ f"(the flattened output-channel axis), got
{weight_axis}"
+ )
+ scale_count =
weight_tensor.tensor.Quantization().ScaleLength()
+ if scale_count != in_channels:
+ raise tvm.error.OpAttributeInvalid(
+ "Per-channel DepthwiseConv2D weight scale count
must match "
+ f"input_channels * depth_multiplier
({in_channels}), got {scale_count}"
+ )
+ # Dequantize while the TFLite [1, KH, KW, C*M] channel
axis is
+ # still intact. Reshaping first would split that axis
into C and M,
+ # which a single-axis dequantize cannot represent.
+ w_f32 = relax.op.dequantize(
+ weight_expr,
+ scale=weight_tensor.qnn_params["scale"],
+ zero_point=weight_tensor.qnn_params["zero_point"],
+ axis=3,
+ )
+ w_f32 = relax.op.reshape(w_f32, depthwise_weight_shape)
+ else:
+ if weight_axis != 0:
+ raise tvm.error.OpAttributeInvalid(
+ "Per-tensor DepthwiseConv2D weight
QuantizedDimension() must be 0, "
+ f"got {weight_axis}"
+ )
+ w_f32 = relax.op.dequantize(
+ weight_expr,
+ scale=weight_tensor.qnn_params["scale"],
+ zero_point=weight_tensor.qnn_params["zero_point"],
+ axis=0,
)
else:
if weight_axis != 0:
@@ -4760,13 +4804,12 @@ class OperatorConverter:
f"Conv2D weight QuantizedDimension() must be 0
(output-channel "
f"axis in [OC,KH,KW,IC] layout), got {weight_axis}"
)
- weight_axis = 3
- w_f32 = relax.op.dequantize(
- weight_expr,
- scale=weight_tensor.qnn_params["scale"],
- zero_point=weight_tensor.qnn_params["zero_point"],
- axis=weight_axis,
- )
+ w_f32 = relax.op.dequantize(
+ weight_expr,
+ scale=weight_tensor.qnn_params["scale"],
+ zero_point=weight_tensor.qnn_params["zero_point"],
+ axis=3,
+ )
# Float convolution
out = relax.op.nn.conv2d(in_f32, w_f32, **params)
else:
@@ -7519,6 +7562,13 @@ class OperatorConverter:
cls_pred = self.get_expr(inputs[1].tensor_idx)
loc_prob = self.get_expr(inputs[0].tensor_idx)
batch_size = inputs[1].tensor.Shape(0)
+ input_num_classes = int(inputs[1].tensor.Shape(2))
+ label_offset = input_num_classes - num_classes
+ if label_offset not in (0, 1):
+ raise ValueError(
+ "DETECTION_POSTPROCESS class predictions must contain
num_classes "
+ "or num_classes + 1 entries"
+ )
anchor_values = self.get_tensor_value(inputs[2])
anchor_boxes = len(anchor_values)
anchor_type = self.get_tensor_type_str(inputs[2].tensor.Type())
@@ -7533,6 +7583,14 @@ class OperatorConverter:
if inputs[2].qnn_params:
anchor_expr = self.dequantize(anchor_expr, inputs[2])
+ if label_offset:
+ cls_pred = relax.op.strided_slice(
+ cls_pred,
+ axes=[2],
+ begin=[label_offset],
+ end=[label_offset + num_classes],
+ )
+
# loc_prob coords are in yxhw format
# need to convert to xywh
loc_coords = relax.op.split(loc_prob, 4, axis=2)
@@ -7569,7 +7627,11 @@ class OperatorConverter:
1 / w_scale,
1 / h_scale,
)
- multibox_transform_loc_attrs["keep_background"] = use_regular_nms
+ # TFLite DetectionPostProcess consumes probabilities, not logits. Any
+ # optional background class was sliced above, so class 0 is now a real
+ # foreground class and must be kept for both NMS implementations.
+ multibox_transform_loc_attrs["keep_background"] = True
+ multibox_transform_loc_attrs["apply_softmax"] = False
multibox_res = self.bb.emit(
relax.op.vision.multibox_transform_loc(
diff --git a/python/tvm/relax/op/vision/multibox_transform_loc.py
b/python/tvm/relax/op/vision/multibox_transform_loc.py
index e4e41a9873..6f5118407a 100644
--- a/python/tvm/relax/op/vision/multibox_transform_loc.py
+++ b/python/tvm/relax/op/vision/multibox_transform_loc.py
@@ -27,8 +27,9 @@ def multibox_transform_loc(
threshold=0.0,
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=True,
+ apply_softmax=True,
):
- """SSD / TFLite-style decode: priors + offsets → boxes; logits → softmax
scores.
+ """SSD / TFLite-style decode: priors + offsets → boxes; prepare class
scores.
Box decode follows TFLite ``DecodeCenterSizeBoxes``; expected tensor
layout matches
``tflite_frontend.convert_detection_postprocess`` (loc reorder yxhw→xywh,
anchor ltrb).
@@ -36,7 +37,7 @@ def multibox_transform_loc(
Parameters
----------
cls_pred : relax.Expr
- ``[B, C, N]`` class logits (pre-softmax).
+ ``[B, C, N]`` class logits or scores.
loc_pred : relax.Expr
``[B, 4*N]`` per-anchor encodings as ``(x,y,w,h)`` after reorder (see
above).
anchor : relax.Expr
@@ -51,6 +52,8 @@ def multibox_transform_loc(
encoded height/width terms inside ``exp(...)`` and can overflow in
float32/float16.
keep_background : bool
If False, set output scores at class index 0 to zero.
+ apply_softmax : bool
+ If True, apply softmax over the class axis before thresholding.
Returns
-------
@@ -82,4 +85,5 @@ def multibox_transform_loc(
threshold,
variances,
keep_background,
+ apply_softmax,
)
diff --git a/python/tvm/relax/transform/legalize_ops/vision.py
b/python/tvm/relax/transform/legalize_ops/vision.py
index 618a30641c..a50ba133f2 100644
--- a/python/tvm/relax/transform/legalize_ops/vision.py
+++ b/python/tvm/relax/transform/legalize_ops/vision.py
@@ -183,6 +183,7 @@ def _multibox_transform_loc(bb: BlockBuilder, call: Call)
-> Expr:
clip=call.attrs.clip,
threshold=call.attrs.threshold,
keep_background=call.attrs.keep_background,
+ apply_softmax=call.attrs.apply_softmax,
)
return bb.call_te(
diff --git a/python/tvm/topi/vision/multibox_transform_loc.py
b/python/tvm/topi/vision/multibox_transform_loc.py
index e6816d8eec..0dc458ebf4 100644
--- a/python/tvm/topi/vision/multibox_transform_loc.py
+++ b/python/tvm/topi/vision/multibox_transform_loc.py
@@ -29,8 +29,9 @@ def multibox_transform_loc(
clip=False,
threshold=0.0,
keep_background=True,
+ apply_softmax=True,
):
- """TFLite ``DecodeCenterSizeBoxes``-style decode + softmax score
post-process.
+ """TFLite ``DecodeCenterSizeBoxes``-style decode + score post-process.
Inputs must match Relax op contracts: ``cls_pred [B,C,N]``, ``loc_pred
[B,4*N]``,
``anchor [1,N,4]`` ltrb; per-anchor loc order ``(x,y,w,h)`` after
yxhw→xywh reorder.
@@ -38,7 +39,7 @@ def multibox_transform_loc(
Parameters
----------
cls_pred : te.Tensor
- ``[B, C, N]`` logits.
+ ``[B, C, N]`` logits or scores.
loc_pred : te.Tensor
``[B, 4*N]`` encodings ``(x,y,w,h)`` per anchor.
anchor : te.Tensor
@@ -48,16 +49,18 @@ def multibox_transform_loc(
clip : bool
Clip ``ymin,xmin,ymax,xmax`` to ``[0,1]``.
threshold : float
- After softmax: ``scores *= (scores >= threshold)``.
+ ``scores *= (scores >= threshold)`` after optional softmax.
keep_background : bool
If False: ``scores[:,0,:] = 0``.
+ apply_softmax : bool
+ Apply softmax across classes before thresholding.
Returns
-------
boxes : te.Tensor
``[B, N, 4]`` as ``(ymin,xmin,ymax,xmax)``.
scores : te.Tensor
- ``[B, C, N]`` softmax, then threshold mask and optional background
zero.
+ ``[B, C, N]`` scores after optional softmax, threshold, and background
masking.
"""
dtype = cls_pred.dtype
B = cls_pred.shape[0]
@@ -107,7 +110,7 @@ def multibox_transform_loc(
boxes = te.compute((B, num_anchors, 4), decode_bbox, name="multibox_boxes")
- scores = topi.nn.softmax(cls_pred, axis=1)
+ scores = topi.nn.softmax(cls_pred, axis=1) if apply_softmax else cls_pred
mask = topi.cast(topi.greater_equal(scores, th), dtype)
scores = scores * mask
if not keep_background:
diff --git a/src/relax/op/vision/multibox_transform_loc.cc
b/src/relax/op/vision/multibox_transform_loc.cc
index 09ac72ea1f..658e382a1e 100644
--- a/src/relax/op/vision/multibox_transform_loc.cc
+++ b/src/relax/op/vision/multibox_transform_loc.cc
@@ -36,7 +36,8 @@ namespace relax {
TVM_FFI_STATIC_INIT_BLOCK() { MultiboxTransformLocAttrs::RegisterReflection();
}
Expr multibox_transform_loc(Expr cls_pred, Expr loc_pred, Expr anchor, bool
clip, double threshold,
- ffi::Array<double> variances, bool
keep_background) {
+ ffi::Array<double> variances, bool keep_background,
+ bool apply_softmax) {
TVM_FFI_ICHECK_EQ(variances.size(), 4)
<< "multibox_transform_loc: variances must be length 4 (x,y,w,h), got "
<< variances.size();
@@ -45,6 +46,7 @@ Expr multibox_transform_loc(Expr cls_pred, Expr loc_pred,
Expr anchor, bool clip
attrs->threshold = threshold;
attrs->variances = std::move(variances);
attrs->keep_background = keep_background;
+ attrs->apply_softmax = apply_softmax;
static const Op& op = Op::Get("relax.vision.multibox_transform_loc");
return Call(Type::Missing(), op, {std::move(cls_pred), std::move(loc_pred),
std::move(anchor)},
@@ -186,12 +188,12 @@ Type InferTypeMultiboxTransformLoc(const Call& call,
const BlockBuilder& ctx) {
TVM_REGISTER_OP("relax.vision.multibox_transform_loc")
.describe(
- "Decode SSD/TFLite-style priors and offsets into boxes and softmax
scores. If "
+ "Decode SSD/TFLite-style priors and offsets into boxes and class
scores. If "
"cls_pred shape is unknown, N-based loc/anchor shape checks are
skipped in "
"inference. Very large variances (w,h) can overflow exp in half box
sizes.")
.set_attrs_type<MultiboxTransformLocAttrs>()
.set_num_inputs(3)
- .add_argument("cls_pred", "Tensor", "[B,C,N] class logits (pre-softmax).")
+ .add_argument("cls_pred", "Tensor", "[B,C,N] class logits or scores.")
.add_argument("loc_pred", "Tensor",
"[B,4*N] box encodings (x,y,w,h); TFLite yxhw order remapped
to xywh.")
.add_argument("anchor", "Tensor", "[1,N,4] priors as ltrb
(left,top,right,bottom).")
diff --git a/src/relax/op/vision/multibox_transform_loc.h
b/src/relax/op/vision/multibox_transform_loc.h
index 726bc4c0e5..7c5a4a77a0 100644
--- a/src/relax/op/vision/multibox_transform_loc.h
+++ b/src/relax/op/vision/multibox_transform_loc.h
@@ -34,7 +34,7 @@ namespace relax {
/*! \brief Decode SSD box encodings and prepare class scores
(TFLite-compatible). */
Expr multibox_transform_loc(Expr cls_pred, Expr loc_pred, Expr anchor, bool
clip, double threshold,
- ffi::Array<double> variances, bool
keep_background);
+ ffi::Array<double> variances, bool
keep_background, bool apply_softmax);
} // namespace relax
} // namespace tvm
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index e7ebbeaf92..54d1c09fe8 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -3577,7 +3577,7 @@ _DETECTION_POSTPROCESS_SMOKE_CASES = [
"num_anchors": 4,
},
2,
- False,
+ True,
id="basic_fast_nms",
),
pytest.param(
@@ -3603,7 +3603,7 @@ _DETECTION_POSTPROCESS_SHAPE_CASES = [
pytest.param(
{
"num_classes": 2,
- "input_num_classes": 5,
+ "input_num_classes": 2,
"max_detections": 2,
"detections_per_class": 2,
"use_regular_nms": False,
@@ -3612,7 +3612,7 @@ _DETECTION_POSTPROCESS_SHAPE_CASES = [
"batch_size": 1,
"num_anchors": 4,
},
- id="wider_input_classes",
+ id="matching_input_classes",
),
pytest.param(
{
@@ -3638,6 +3638,23 @@ _DETECTION_POSTPROCESS_SHAPE_CASES = [
def test_detection_postprocess_smoke(build_kwargs, expected_topk_count,
expected_keep_background):
mod = _build_detection_postprocess_mod(**build_kwargs)
+ topk_calls = []
+ multibox_calls = []
+
+ def _visit(expr):
+ if isinstance(expr, relax.Call) and expr.op ==
tvm.ir.Op.get("relax.topk"):
+ topk_calls.append(expr)
+ if isinstance(expr, relax.Call) and expr.op == tvm.ir.Op.get(
+ "relax.vision.multibox_transform_loc"
+ ):
+ multibox_calls.append(expr)
+
+ relax.analysis.post_order_visit(mod["main"].body, _visit)
+ assert len(topk_calls) == expected_topk_count
+ assert len(multibox_calls) == 1
+ assert multibox_calls[0].attrs.keep_background == expected_keep_background
+ assert not multibox_calls[0].attrs.apply_softmax
+
expected_batch = build_kwargs["batch_size"]
expected_max_detections = build_kwargs["max_detections"]
tvm.ir.assert_structural_equal(
@@ -3681,6 +3698,38 @@ def
test_detection_postprocess_shape_variations(build_kwargs):
)
+def test_detection_postprocess_removes_background_without_softmax():
+ """TFLite scores are probabilities; remove its optional background class
exactly once."""
+ mod = _build_detection_postprocess_mod(
+ num_classes=2,
+ input_num_classes=3,
+ max_detections=2,
+ detections_per_class=2,
+ batch_size=1,
+ )
+ multibox_calls = []
+
+ def _visit(expr):
+ if isinstance(expr, relax.Call) and expr.op == tvm.ir.Op.get(
+ "relax.vision.multibox_transform_loc"
+ ):
+ multibox_calls.append(expr)
+
+ relax.analysis.post_order_visit(mod["main"].body, _visit)
+ assert len(multibox_calls) == 1
+ assert not multibox_calls[0].attrs.apply_softmax
+ assert multibox_calls[0].attrs.keep_background
+ tvm.ir.assert_structural_equal(
+ multibox_calls[0].args[0].ty,
+ relax.TensorType((1, 2, 4), "float32"),
+ )
+
+
+def test_detection_postprocess_rejects_invalid_class_count():
+ with pytest.raises(ValueError, match=r"num_classes \+ 1"):
+ _build_detection_postprocess_mod(num_classes=2, input_num_classes=5)
+
+
def _make_resize_expected(
input_shape, output_size, method, coordinate_transformation_mode,
rounding_method
):
@@ -11352,6 +11401,10 @@ def test_quantized_avg_pool2d_uses_astype():
else:
tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+ # Exercise the public entry point so quantized pool allowlist regressions
+ # cannot be hidden by calling the converter method directly.
+ from_tflite(tflite_model)
+
subgraph = tflite_model.Subgraphs(0)
bb = relax.BlockBuilder()
exp_tab = tflite_frontend.ExprTable()
@@ -12315,8 +12368,8 @@ def
test_quantized_conv2d_per_channel_weight_with_int32_bias_dequantizes_bias():
tvm.ir.assert_structural_equal(mod, Expected)
-def test_per_channel_depthwise_conv_unsupported():
- """Per-channel quantized depthwise Conv2D raises OpNotImplemented."""
+def test_per_channel_depthwise_conv_dequantizes_before_reshape():
+ """Per-channel depthwise weights keep C*M intact and lower to HWOI."""
import flatbuffers
import tflite.Model
@@ -12325,9 +12378,12 @@ def test_per_channel_depthwise_conv_unsupported():
in_q = _build_quantization_parameters(
builder, scale=[0.5], zero_point=[0], quantized_dimension=0
)
- # Per-channel weight: 2 channels, scale vector length 2
+ # Two input channels with depth_multiplier=2 produce four output channels.
wt_q = _build_quantization_parameters(
- builder, scale=[0.25, 0.75], zero_point=[0, 0], quantized_dimension=3
+ builder,
+ scale=[0.25, 0.5, 0.75, 1.0],
+ zero_point=[0, 0, 0, 0],
+ quantized_dimension=3,
)
out_q = _build_quantization_parameters(
builder, scale=[1.0], zero_point=[0], quantized_dimension=0
@@ -12337,16 +12393,16 @@ def test_per_channel_depthwise_conv_unsupported():
builder, 0, [1, 4, 4, 2], tensor_type=_tfl_tensor_type.INT8,
quantization=in_q
)
t_wt = _build_tensor(
- builder, 1, [1, 3, 3, 2], tensor_type=_tfl_tensor_type.INT8,
quantization=wt_q
+ builder, 1, [1, 3, 3, 4], tensor_type=_tfl_tensor_type.INT8,
quantization=wt_q
)
t_ou = _build_tensor(
- builder, 2, [1, 2, 2, 2], tensor_type=_tfl_tensor_type.INT8,
quantization=out_q
+ builder, 2, [1, 2, 2, 4], tensor_type=_tfl_tensor_type.INT8,
quantization=out_q
)
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsStart(builder)
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddStrideH(builder, 1)
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddStrideW(builder, 1)
-
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddDepthMultiplier(builder,
1)
+
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddDepthMultiplier(builder,
2)
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddPadding(builder, 1)
_tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsAddFusedActivationFunction(builder,
0)
dw_opts = _tfl_depthwise_conv2d_options.DepthwiseConv2DOptionsEnd(builder)
@@ -12379,8 +12435,25 @@ def test_per_channel_depthwise_conv_unsupported():
else:
tflite_model = tflite.Model.GetRootAsModel(buf, 0)
- with pytest.raises(tvm.error.OpNotImplemented, match="Per-channel"):
- from_tflite(tflite_model)
+ mod = from_tflite(tflite_model)
+ dequantize_calls = []
+ reshape_calls = []
+
+ def _visit(expr):
+ if isinstance(expr, relax.Call) and expr.op ==
tvm.ir.Op.get("relax.dequantize"):
+ dequantize_calls.append(expr)
+ if isinstance(expr, relax.Call) and expr.op ==
tvm.ir.Op.get("relax.reshape"):
+ reshape_calls.append(expr)
+
+ relax.analysis.post_order_visit(mod["main"].body, _visit)
+ assert any(call.attrs.axis == 3 for call in dequantize_calls)
+ depthwise_reshapes = [
+ call
+ for call in reshape_calls
+ if call.ty.dtype == "float32"
+ and tuple(dim.value for dim in call.ty.shape.values) == (3, 3, 4, 1)
+ ]
+ assert len(depthwise_reshapes) == 1
def test_uint8_reshape_requantize_uses_dq_reshape_q():
diff --git a/tests/python/relax/test_op_vision.py
b/tests/python/relax/test_op_vision.py
index 9010f74606..abdbb6a42e 100644
--- a/tests/python/relax/test_op_vision.py
+++ b/tests/python/relax/test_op_vision.py
@@ -1490,7 +1490,14 @@ def test_multibox_transform_loc_wrong_batch():
def _multibox_ref_numpy(
- cls_pred, loc_pred, anchor, variances, clip=False, threshold=0.0,
keep_background=True
+ cls_pred,
+ loc_pred,
+ anchor,
+ variances,
+ clip=False,
+ threshold=0.0,
+ keep_background=True,
+ apply_softmax=True,
):
"""Numpy reference aligned with ``topi.vision.multibox_transform_loc``."""
@@ -1501,7 +1508,11 @@ def _multibox_ref_numpy(
B, C, N = cls_pred.shape
loc = loc_pred.reshape(B, N, 4)
- scores = _softmax(cls_pred.astype("float64"), axis=1).astype(np.float32)
+ scores = (
+ _softmax(cls_pred.astype("float64"), axis=1).astype(np.float32)
+ if apply_softmax
+ else cls_pred.copy()
+ )
if threshold > 0.0:
scores = np.where(scores >= threshold, scores, 0.0).astype(np.float32)
if not keep_background:
@@ -1648,6 +1659,7 @@ def test_multibox_transform_loc_legalize_attr_branches():
threshold=0.4,
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=False,
+ apply_softmax=False,
)
cls_data = np.array(
@@ -1674,6 +1686,7 @@ def test_multibox_transform_loc_legalize_attr_branches():
clip=True,
threshold=0.4,
keep_background=False,
+ apply_softmax=False,
)
out = vm["main"](
tvm.runtime.tensor(cls_data, tvm.cpu()),
diff --git a/tests/python/relax/test_tvmscript_parser_op_vision.py
b/tests/python/relax/test_tvmscript_parser_op_vision.py
index ac5fa78b39..8ab952adf5 100644
--- a/tests/python/relax/test_tvmscript_parser_op_vision.py
+++ b/tests/python/relax/test_tvmscript_parser_op_vision.py
@@ -291,6 +291,7 @@ def test_multibox_transform_loc():
threshold=0.0,
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=True,
+ apply_softmax=True,
)
)
return gv
@@ -310,6 +311,7 @@ def test_multibox_transform_loc():
threshold=0.0,
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=True,
+ apply_softmax=True,
)
)
bb.emit_func_output(gv)
