diff --git a/ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py b/ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py index e2276d7..a04f32a 100644 --- a/ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py +++ b/ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py @@ -28,6 +28,9 @@ def check_op_quantization_config( Raises: ValueError: If the op quantization config is invalid. """ + if op_quant_config.skip_checks: + return + if op_quant_config.weight_tensor_config.dtype != qtyping.TensorDataType.INT: raise ValueError( "Weights need to have integer type for min/max uniform quantization. If" diff --git a/ai_edge_quantizer/qtyping.py b/ai_edge_quantizer/qtyping.py index ba8d6b8..94f8e95 100644 --- a/ai_edge_quantizer/qtyping.py +++ b/ai_edge_quantizer/qtyping.py @@ -246,6 +246,7 @@ class OpQuantizationConfig: weight_tensor_config: The quantization configuration for weight tensor in the op. execution_mode: How to execute the op after quantization. + skip_checks: Skip op quantization config checks. """ # Quant config for activation tensors in the op (i.e., runtime tensors). @@ -262,6 +263,11 @@ class OpQuantizationConfig: # How to execute the op after quantization. execution_mode: OpExecutionMode = OpExecutionMode.WEIGHT_ONLY + # For advanced users only. If set, the quantizer will ignore all op + # configuration checks and forcefully quantize this op according to the user + # instructions even if it's not supported in the TFLite runtime. + skip_checks: bool = False + def __post_init__(self): if self.activation_tensor_config is None: return