TinyNeuralNetwork
TinyNeuralNetwork copied to clipboard
Model with stack does not work with int8 target type
Converting this dummy model with quantize_target_type="int8" and per_tensor=True throws an error in tflite
import torch.nn as nn
import torch
from tinynn.graph.quantization.quantizer import PostQuantizer
from tinynn.converter import TFLiteConverter
class StackModel(nn.Module):
def forward(self, x: torch.Tensor):
"""
Args:
x: [N, H, W, C]
"""
return torch.stack([-x, x], dim=-1)
def _main():
dummy_input = torch.rand(1, 60, 60, 256).float()
model = StackModel()
qat_config = {
"backend": "qnnpack",
"per_tensor": True,
"disable_requantization_for_cat": True
}
quantizer = PostQuantizer(
model, (dummy_input,), work_dir="stack_model", config=qat_config
)
ptq_coarse_matcher = quantizer.quantize()
ptq_coarse_matcher(dummy_input)
with torch.no_grad():
ptq_coarse_matcher.eval()
ptq_coarse_matcher.cpu()
ptq_coarse_matcher = quantizer.convert(ptq_coarse_matcher)
torch.backends.quantized.engine = quantizer.backend
converter = TFLiteConverter(
ptq_coarse_matcher,
(dummy_input),
"stack_model.tflite",
fuse_quant_dequant=True,
quantize_target_type="int8"
)
converter.convert()
if __name__ == '__main__':
_main()
Tflite error:
return self._interpreter.AllocateTensors()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: /tensorflow/tensorflow/lite/kernels/concatenation.cc:184 t->params.zero_point != output->params.zero_point (-1 != 0)Node number 3 (CONCATENATION) failed to prepare.
Note that the model works fine if I remove the "negative x" and instead send the same tensor twice, and it works with uint8