diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 13d78c4c30a5..fe6c420d5bae 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -101,6 +101,7 @@ class OperatorConverter: "ABS", "ADD", "ATAN2", + "AVERAGE_POOL_2D", "CEIL", "CONCATENATION", "CONV_2D", @@ -5640,6 +5641,50 @@ def convert_topk_v2(self, op): return out + def _pool2d_static_hw(self, expr, what): + """(H, W) of an NHWC Relax tensor, which must be static. + + Pool2D's SAME padding and the quantized average pool's divisor are + folded to constants, so they need concrete extents -- read from the + Relax expression rather than the serialized TFLite tensor, which is + stale once shape_dict overrides the input shape.""" + shape = self._infer_shape(expr) + try: + _, h, w, _ = [int(v) for v in shape] + except (TypeError, ValueError) as err: + raise tvm.error.OpAttributeUnImplemented( + f"Pool2D requires a static NHWC {what} shape, got {shape}" + ) from err + return h, w + + @staticmethod + def _avg_pool2d_valid_counts(in_hw, filter_hw, stride_hw, padding, out_hw): + """Non-padded taps per output position, as an int32 [1, OH, OW, 1] array. + + TFLite's average pool divides by the number of taps that fall inside + the input (`count_include_pad=False` semantics). With no padding that + is just filter_h*filter_w everywhere; with SAME padding the windows at + the borders see fewer, so the divisor is position dependent. + """ + import numpy as _np + + in_h, in_w = in_hw + f_h, f_w = filter_hw + s_h, s_w = stride_hw + pad_top, pad_left = (padding[0], padding[1]) if len(padding) >= 2 else (0, 0) + out_h, out_w = out_hw + + def counts_1d(extent, f, s, pad_before, out): + c = _np.empty(out, dtype="int32") + for i in range(out): + start = i * s - pad_before + c[i] = min(start + f, extent) - max(start, 0) + return c + + rows = counts_1d(in_h, f_h, s_h, pad_top, out_h) + cols = counts_1d(in_w, f_w, s_w, pad_left, out_w) + return (rows[:, None] * cols[None, :]).reshape(1, out_h, out_w, 1) + def convert_pool2d(self, op, pool_type): """pool2d implementation.""" @@ -5678,7 +5723,11 @@ def convert_pool2d(self, op, pool_type): in_expr = self.get_expr(input_tensor_idx) - _, input_h, input_w, _ = to_int_list(self.get_tensor_shape(input_tensor)) + # Take H and W from the Relax input, not from the serialized TFLite + # shape: from_tflite(..., shape_dict=...) can override the input + # dimensions, and then the flatbuffer's shapes are stale. SAME padding + # and the quantized divisor below both depend on the real extents. + input_h, input_w = self._pool2d_static_hw(in_expr, "input") if padding == Padding.VALID: pass @@ -5697,8 +5746,78 @@ def convert_pool2d(self, op, pool_type): "TFLite avg_pool2dreshape requires input and output scale" "and zero points to be equal" ) - out = relax.op.astype(in_expr, "int32") - out = relax.op.nn.avg_pool2d(out, **params) + # TFLite's quantized AveragePool rounds the window average + # HALF AWAY FROM ZERO, and relax's integer avg_pool2d divides + # with a truncating division. Truncating biases every pooled + # value toward zero by up to half an LSB -- on a quantized + # ResNet that is ~52% of the pooled values, and because such + # graphs end in an int8 softmax it reaches the logits as tens + # of counts. The reference is + # tensorflow/lite/kernels/internal/reference/integer_ops/ + # pooling.h + # acc = acc > 0 ? (acc + count / 2) / count + # : (acc - count / 2) / count + # + # So take the window SUM and do that division explicitly. + # avg_pool2d divides by the window size, so pre-scaling the + # input by that size makes its division exact and what comes + # back is the sum. count_include_pad=True keeps that divisor a + # constant; the padded taps are zeros and contribute nothing, + # so the result is the sum over the VALID taps either way. + window = filter_h * filter_w + # The accumulator must hold the PRE-SCALED window sum. Every + # element is scaled by `window` and then `window` of them are + # summed, so |acc| <= max|x| * window^2 -- and max|x| is set by + # the INPUT type, which is not always 8-bit: this branch also + # takes int16 (a 17x17 pool of int16 32767 needs 2.7e9, past + # int32). Size the accumulator from the actual bound: int32 + # whenever it fits, which keeps the usual int8 case as cheap as + # before, else int64. + in_info = np.iinfo(self.get_tensor_type_str(input_tensor.tensor.Type())) + max_abs = max(-int(in_info.min), int(in_info.max)) # 128 int8, 255 uint8 + bound = max_abs * window * window + if bound < (1 << 31): + acc_dtype = "int32" + elif bound < (1 << 63): + acc_dtype = "int64" + else: + raise tvm.error.OpAttributeUnImplemented( + f"pooling window {filter_h}x{filter_w} is too large for an " + "exact-sum average pool even with an int64 accumulator" + ) + acc = relax.op.astype(in_expr, acc_dtype) + acc = relax.op.multiply(acc, relax.const(window, acc_dtype)) + acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) + # relax widens the output dtype of an INTEGER avg_pool2d when the + # window is large enough that an int32 accumulator could + # overflow -- an 8x8 pool over int32 comes back as int64, a 2x2 + # one stays int32. Pin it back to the accumulator type so the + # rounding arithmetic below always meets constants of the same + # dtype; without this an 8x8 pool fails to import with "Binary + # operators must have the same datatype". + acc = self.bb.normalize(relax.op.astype(acc, acc_dtype)) + + # TFLite divides by the number of NON-PADDED taps, which varies + # per output position once there is padding. The shapes are + # static, so the per-position count is folded to a constant + # here rather than computed in the graph. Both extents come from + # the Relax graph -- the input above and the pooled output here + # -- so the constant matches the tensor it divides even when + # shape_dict overrides the serialized input shape. + counts = self._avg_pool2d_valid_counts( + (input_h, input_w), + (filter_h, filter_w), + (stride_h, stride_w), + params["padding"], + self._pool2d_static_hw(acc, "pooled output"), + ) + half = relax.const(counts // 2, acc_dtype) + out = relax.op.where( + relax.op.greater(acc, relax.const(0, acc_dtype)), + relax.op.add(acc, half), + relax.op.subtract(acc, half), + ) + out = relax.op.divide(out, relax.const(counts, acc_dtype)) out = relax.op.astype(out, output_tensor_type_str) else: out = relax.op.nn.avg_pool2d(in_expr, **params) diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index e7ebbeaf92ce..14c14edeb5d7 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -57,7 +57,6 @@ def _get_mod_from_cfunc(cfunc): mod["main"] = mod["main"].without_attr("params") return mod - def verify(TestClass, expected=None): if isinstance(TestClass, type): cf = TestClass().func.get_concrete_function() @@ -11293,8 +11292,8 @@ def main() -> R.Tensor((2,), dtype="float32"): tvm.ir.assert_structural_equal(mod, Expected) -def test_quantized_avg_pool2d_uses_astype(): - """Quantized AVERAGE_POOL_2D casts through int32 with R.astype.""" +def test_quantized_avg_pool2d_rounds_half_away_from_zero(): + """Quantized AVERAGE_POOL_2D takes the window sum and rounds like TFLite.""" builder = flatbuffers.Builder(1024) qparams = _build_quantization_parameters( @@ -11373,24 +11372,314 @@ def main(tvmgen_tensor_0: R.Tensor((1, 2, 2, 1), dtype="int8")) -> R.Tensor( ): with R.dataflow(): lv: R.Tensor((1, 2, 2, 1), dtype="int32") = R.astype(tvmgen_tensor_0, dtype="int32") - lv1: R.Tensor((1, 1, 1, 1), dtype="int32") = R.nn.avg_pool2d( - lv, + # scaling by the window size makes avg_pool2d's own division + # exact, so what comes back is the window SUM + lv1: R.Tensor((1, 2, 2, 1), dtype="int32") = R.multiply(lv, R.const(4, "int32")) + lv2: R.Tensor((1, 1, 1, 1), dtype="int32") = R.nn.avg_pool2d( + lv1, pool_size=[2, 2], strides=[1, 1], dilation=[1, 1], padding=[0, 0, 0, 0], ceil_mode=False, - count_include_pad=False, + count_include_pad=True, layout="NHWC", out_layout="NHWC", ) - gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv1, dtype="int8") + lv3: R.Tensor((1, 1, 1, 1), dtype="int32") = R.astype(lv2, dtype="int32") + # acc > 0 ? (acc + count/2) / count : (acc - count/2) / count + lv4: R.Tensor((1, 1, 1, 1), dtype="bool") = R.greater(lv3, R.const(0, "int32")) + lv5: R.Tensor((1, 1, 1, 1), dtype="int32") = R.add( + lv3, R.const(np.full((1, 1, 1, 1), 2, dtype="int32")) + ) + lv6: R.Tensor((1, 1, 1, 1), dtype="int32") = R.subtract( + lv3, R.const(np.full((1, 1, 1, 1), 2, dtype="int32")) + ) + lv7: R.Tensor((1, 1, 1, 1), dtype="int32") = R.where(lv4, lv5, lv6) + lv8: R.Tensor((1, 1, 1, 1), dtype="int32") = R.divide( + lv7, R.const(np.full((1, 1, 1, 1), 4, dtype="int32")) + ) + gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv8, dtype="int8") R.output(gv) return gv tvm.ir.assert_structural_equal(mod, Expected) +def test_quantized_avg_pool2d_matches_tflite_rounding_numerically(): + """The pooled values must match TFLite's reference integer average pool. + + The structural test above pins the shape of the lowering; this one pins the + ANSWER, which is what actually regressed: a truncating division looks + reasonable and is wrong on about half of all inputs. The reference is + tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h, + + acc = acc > 0 ? (acc + count / 2) / count : (acc - count / 2) / count + + and the inputs below are chosen so the window sums land on and around the + .5 boundaries where truncation and round-half-away disagree. + + The activation was also tested with resnet. + """ + builder = flatbuffers.Builder(1024) + + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, 2, 2, 8], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 1, 1, 8], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, _tfl_padding.VALID) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + mod = from_tflite(tflite_model) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + + # one channel per window sum in [-7, 7]: +-1.5, +-1.75 and the exact halves + # are where a truncating divide diverges from round-half-away + sums = [5, 6, 7, -5, -6, -7, 2, -2] + x = np.zeros((1, 2, 2, 8), dtype="int8") + for c, total in enumerate(sums): + x[0, 0, 0, c] = total + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy().reshape(-1) + + count = 4 + want = np.array( + [(t + count // 2) // count if t > 0 else -((-t + count // 2) // count) for t in sums], + dtype="int8", + ) + np.testing.assert_array_equal(got, want) + + +@pytest.mark.parametrize( + "tensor_type, dtype, value", + [ + # The pre-scaled window sum is max|x| * window^2: for int16 17x17 that is + # 2.7e9, past int32, and an int32 accumulator returned -18657 / 18656. + (_tfl_tensor_type.INT16, "int16", 32767), + (_tfl_tensor_type.INT16, "int16", -32768), + (_tfl_tensor_type.INT8, "int8", 127), + (_tfl_tensor_type.INT8, "int8", -128), + (_tfl_tensor_type.UINT8, "uint8", 255), + ], +) +def test_quantized_avg_pool2d_large_window_does_not_overflow(tensor_type, dtype, value): + """A 17x17 VALID average pool of a constant input must return that constant. + + The lowering computes the exact window sum by pre-scaling the input by the + window size, so its accumulator has to hold max|x| * window^2 -- which + depends on the INPUT type, not just the window. A constant input makes the + answer obvious (the average of N copies of v is v) and puts the sum at its + extreme. + """ + size = 17 + builder = flatbuffers.Builder(1024) + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, size, size, 1], tensor_type=tensor_type, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 1, 1, 1], tensor_type=tensor_type, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, _tfl_padding.VALID) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 1) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 1) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, size) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, size) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + mod = from_tflite(tflite_model) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + x = np.full((1, size, size, 1), value, dtype=dtype) + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy().reshape(-1) + np.testing.assert_array_equal(got, np.array([value], dtype=dtype)) + + +def _tflite_int_avg_pool2d_reference(x, filter_hw, stride_hw, same_padding): + """TFLite's reference integer AveragePool on an NHWC array: divide the sum over + the NON-PADDED taps, rounding half away from zero (reference/integer_ops/pooling.h).""" + n, in_h, in_w, c = x.shape + (f_h, f_w), (s_h, s_w) = filter_hw, stride_hw + + def axis(extent, f, s): + if not same_padding: + return (extent - f) // s + 1, 0 + out = -(-extent // s) + return out, max((out - 1) * s + f - extent, 0) // 2 + + out_h, pad_h = axis(in_h, f_h, s_h) + out_w, pad_w = axis(in_w, f_w, s_w) + y = np.zeros((n, out_h, out_w, c), dtype=x.dtype) + for oy in range(out_h): + for ox in range(out_w): + y0, x0 = oy * s_h - pad_h, ox * s_w - pad_w + win = x[:, max(y0, 0) : min(y0 + f_h, in_h), max(x0, 0) : min(x0 + f_w, in_w), :] + count = win.shape[1] * win.shape[2] + acc = win.astype("int64").sum(axis=(1, 2)) + y[:, oy, ox, :] = np.where( + acc > 0, (acc + count // 2) // count, -((-acc + count // 2) // count) + ) + return y + + +@pytest.mark.parametrize( + "padding, override, want_hw", + [ + # the review's cases: the serialized model is 4x4 -> 2x2 + (_tfl_padding.VALID, (1, 2, 2, 1), (1, 1)), # counts tensor used to broadcast to 2x2 + (_tfl_padding.VALID, (1, 6, 6, 1), (3, 3)), # used to fail: 3x3 vs a 2x2 counts tensor + # SAME: the border windows see fewer taps, and the padding itself must be + # derived from the overridden extent, not the serialized 4 + (_tfl_padding.SAME, (1, 5, 5, 1), (3, 3)), + (_tfl_padding.SAME, (1, 7, 3, 1), (4, 2)), + ], +) +def test_quantized_avg_pool2d_follows_shape_dict_override(padding, override, want_hw): + """The quantized AVERAGE_POOL_2D divisor must come from the Relax shapes. + + The per-position counts (and SAME padding) are folded to constants. They used + to be built from the SERIALIZED TFLite shapes, which are stale once + from_tflite(..., shape_dict=...) overrides the input: a smaller input got a + counts tensor that silently broadcast the result back to the old output + shape, and a larger one failed to import. + """ + builder = flatbuffers.Builder(1024) + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, 4, 4, 1], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 2, 2, 1], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, padding) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + from tvm.relax.frontend.tflite.tflite_frontend import _input_type + + input_name = next(iter(_input_type(tflite_model)[0])) + mod = from_tflite(tflite_model, shape_dict={input_name: override}) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + x = np.random.default_rng(0).integers(-128, 128, size=override, dtype=np.int64).astype("int8") + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy() + + assert got.shape == (1, *want_hw, 1) + want = _tflite_int_avg_pool2d_reference( + x, (2, 2), (2, 2), same_padding=padding == _tfl_padding.SAME + ) + np.testing.assert_array_equal(got, want) + + def test_quantized_conv2d_per_tensor_uses_qdq(): """Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q.""" builder = flatbuffers.Builder(2048)