1# Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Tests for QuantizationDebugger.""" 16 17import csv 18import io 19import re 20 21from unittest import mock 22from absl.testing import parameterized 23 24import numpy as np 25import tensorflow as tf 26 27from tensorflow.lite.python import convert 28from tensorflow.lite.python import lite 29from tensorflow.lite.tools.optimize.debugging.python import debugger 30from tensorflow.python.framework import test_util 31from tensorflow.python.platform import test 32from tensorflow.python.training.tracking import tracking 33 34# pylint: disable=g-import-not-at-top 35try: 36 from tensorflow.lite.python import metrics_portable as metrics 37except ImportError: 38 from tensorflow.lite.python import metrics_nonportable as metrics 39# pylint: enable=g-import-not-at-top 40 41 42def _get_model(): 43 """Returns somple model with Conv2D and representative dataset gen.""" 44 root = tracking.AutoTrackable() 45 kernel_in = np.array([-2, -1, 1, 2], dtype=np.float32).reshape((2, 2, 1, 1)) 46 47 @tf.function( 48 input_signature=[tf.TensorSpec(shape=[1, 3, 3, 1], dtype=tf.float32)]) 49 def func(inp): 50 kernel = tf.constant(kernel_in, dtype=tf.float32) 51 conv = tf.nn.conv2d(inp, kernel, strides=1, padding='SAME') 52 output = tf.nn.relu(conv, name='output') 53 return output 54 55 root.f = func 56 to_save = root.f.get_concrete_function() 57 return (root, to_save) 58 59 60def _calibration_gen(): 61 for i in range(5): 62 yield [np.arange(9).reshape((1, 3, 3, 1)).astype(np.float32) * i] 63 64 65def _convert_model(model, func): 66 """Converts TF model to TFLite float model.""" 67 converter = lite.TFLiteConverterV2.from_concrete_functions([func], model) 68 # TODO(b/191205988): Explicitly disable saved model lowering in conversion. 69 converter.experimental_lower_to_saved_model = False 70 return converter.convert() 71 72 73def _quantize_converter(model, func, calibration_gen, debug=True): 74 """Returns a converter appropriate for the function and debug configs.""" 75 converter = lite.TFLiteConverterV2.from_concrete_functions([func], model) 76 converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8] 77 converter.representative_dataset = calibration_gen 78 79 # TODO(b/191205988): Explicitly disable saved model lowering in conversion. 80 converter.experimental_lower_to_saved_model = False 81 82 # Create a TFLite model with new quantizer and numeric verify ops. 83 converter.optimizations = [lite.Optimize.DEFAULT] 84 converter.experimental_new_quantizer = True 85 if debug: 86 converter._experimental_calibrate_only = True 87 return converter 88 89 90def _quantize_model(model, 91 func, 92 calibration_gen, 93 quantized_io=False, 94 debug=True): 95 """Quantizes model, in debug or normal mode.""" 96 converter = _quantize_converter(model, func, calibration_gen, debug) 97 if debug: 98 calibrated = converter.convert() 99 return convert.mlir_quantize( 100 calibrated, enable_numeric_verify=True, fully_quantize=quantized_io) 101 else: 102 return converter.convert() 103 104 105def _dummy_fn(*unused_args): 106 return 0.0 107 108 109class QuantizationDebugOptionsTest(test_util.TensorFlowTestCase, 110 parameterized.TestCase): 111 112 @test_util.run_v2_only 113 def test_init_duplicate_keys_raises_ValueError(self): 114 with self.assertRaises(ValueError): 115 debugger.QuantizationDebugOptions( 116 layer_debug_metrics={ 117 'a': _dummy_fn, 118 'b': _dummy_fn 119 }, 120 model_debug_metrics={ 121 'c': _dummy_fn, 122 'd': _dummy_fn 123 }, 124 layer_direct_compare_metrics={ 125 'a': _dummy_fn, 126 'e': _dummy_fn 127 }) 128 129 with self.assertRaises(ValueError): 130 debugger.QuantizationDebugOptions( 131 layer_debug_metrics={ 132 'a': _dummy_fn, 133 'b': _dummy_fn 134 }, 135 layer_direct_compare_metrics={ 136 'a': _dummy_fn, 137 'e': _dummy_fn 138 }) 139 140 141class QuantizationDebuggerTest(test_util.TensorFlowTestCase, 142 parameterized.TestCase): 143 144 @classmethod 145 def setUpClass(cls): 146 super().setUpClass() 147 cls.tf_model_root, cls.tf_model = _get_model() 148 cls.float_model = _convert_model(cls.tf_model_root, cls.tf_model) 149 cls.debug_model_float = _quantize_model( 150 cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=False) 151 cls.debug_model_int8 = _quantize_model( 152 cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=True) 153 154 @parameterized.named_parameters( 155 ('float_io', False, False), 156 ('quantized_io', True, False), 157 ('float_io_from_converter', False, True), 158 ('quantized_io_from_converter', True, True), 159 ) 160 @test_util.run_v2_only 161 def test_layer_metrics(self, quantized_io, from_converter): 162 options = debugger.QuantizationDebugOptions( 163 layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))}) 164 if not from_converter: 165 if quantized_io: 166 debug_model = QuantizationDebuggerTest.debug_model_int8 167 else: 168 debug_model = QuantizationDebuggerTest.debug_model_float 169 quant_debugger = debugger.QuantizationDebugger( 170 quant_debug_model_content=debug_model, 171 debug_dataset=_calibration_gen, 172 debug_options=options) 173 else: 174 options.fully_quantize = quantized_io 175 quant_debugger = debugger.QuantizationDebugger( 176 converter=_quantize_converter(self.tf_model_root, self.tf_model, 177 _calibration_gen), 178 debug_dataset=_calibration_gen, 179 debug_options=options) 180 181 quant_debugger.run() 182 183 expected_metrics = { 184 'num_elements': 9, 185 'stddev': 0.03850026, 186 'mean_error': 0.01673192, 187 'max_abs_error': 0.10039272, 188 'mean_squared_error': 0.0027558778, 189 'l1_norm': 0.023704167, 190 } 191 self.assertLen(quant_debugger.layer_statistics, 1) 192 actual_metrics = next(iter(quant_debugger.layer_statistics.values())) 193 194 self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys()) 195 for key, value in expected_metrics.items(): 196 self.assertAlmostEqual(value, actual_metrics[key], places=5) 197 198 buffer = io.StringIO() 199 quant_debugger.layer_statistics_dump(buffer) 200 reader = csv.DictReader(buffer.getvalue().split()) 201 actual_values = next(iter(reader)) 202 203 expected_values = expected_metrics.copy() 204 expected_values.update({ 205 'op_name': 'CONV_2D', 206 'tensor_idx': 7 if quantized_io else 8, 207 'scale': 0.15686275, 208 'zero_point': -128, 209 'tensor_name': r'Identity[1-9]?$' 210 }) 211 for key, value in expected_values.items(): 212 if isinstance(value, str): 213 self.assertIsNotNone( 214 re.match(value, actual_values[key]), 215 'String is different from expected string. Please fix test code if' 216 " it's being affected by graph manipulation changes.") 217 elif isinstance(value, list): 218 self.assertAlmostEqual( 219 value[0], float(actual_values[key][1:-1]), places=5) 220 else: 221 self.assertAlmostEqual(value, float(actual_values[key]), places=5) 222 223 @parameterized.named_parameters( 224 ('float_io', False), 225 ('quantized_io', True), 226 ) 227 @test_util.run_v2_only 228 def test_model_metrics(self, quantized_io): 229 if quantized_io: 230 debug_model = QuantizationDebuggerTest.debug_model_int8 231 else: 232 debug_model = QuantizationDebuggerTest.debug_model_float 233 options = debugger.QuantizationDebugOptions( 234 model_debug_metrics={'stdev': lambda x, y: np.std(x[0] - y[0])}) 235 quant_debugger = debugger.QuantizationDebugger( 236 quant_debug_model_content=debug_model, 237 float_model_content=QuantizationDebuggerTest.float_model, 238 debug_dataset=_calibration_gen, 239 debug_options=options) 240 quant_debugger.run() 241 242 expected_metrics = {'stdev': 0.050998904} 243 actual_metrics = quant_debugger.model_statistics 244 245 self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys()) 246 for key, value in expected_metrics.items(): 247 self.assertAlmostEqual(value, actual_metrics[key], places=5) 248 249 @parameterized.named_parameters( 250 ('float_io', False), 251 ('quantized_io', True), 252 ) 253 @test_util.run_v2_only 254 def test_layer_direct_compare_metrics(self, quantized_io): 255 def _corr(float_values, quant_values, scale, zero_point): 256 dequant_values = (quant_values.astype(np.int32) - zero_point) * scale 257 return np.corrcoef(float_values.flatten(), dequant_values.flatten())[0, 1] 258 259 if quantized_io: 260 debug_model = QuantizationDebuggerTest.debug_model_int8 261 else: 262 debug_model = QuantizationDebuggerTest.debug_model_float 263 264 options = debugger.QuantizationDebugOptions( 265 layer_direct_compare_metrics={'corr': _corr}) 266 quant_debugger = debugger.QuantizationDebugger( 267 quant_debug_model_content=debug_model, 268 debug_dataset=_calibration_gen, 269 debug_options=options) 270 quant_debugger.run() 271 272 expected_metrics = { 273 'corr': 0.99999, 274 } 275 self.assertLen(quant_debugger.layer_statistics, 1) 276 actual_metrics = next(iter(quant_debugger.layer_statistics.values())) 277 278 for key, value in expected_metrics.items(): 279 self.assertAlmostEqual(value, actual_metrics[key], places=5) 280 281 @test_util.run_v2_only 282 def test_wrong_input_raises_ValueError(self): 283 284 def wrong_calibration_gen(): 285 for _ in range(5): 286 yield [ 287 np.ones((1, 3, 3, 1), dtype=np.float32), 288 np.ones((1, 3, 3, 1), dtype=np.float32) 289 ] 290 291 quant_debugger = debugger.QuantizationDebugger( 292 quant_debug_model_content=QuantizationDebuggerTest.debug_model_float, 293 debug_dataset=wrong_calibration_gen) 294 with self.assertRaisesRegex( 295 ValueError, r'inputs provided \(2\).+inputs to the model \(1\)'): 296 quant_debugger.run() 297 298 @test_util.run_v2_only 299 def test_non_debug_model_raises_ValueError(self): 300 normal_quant_model = _quantize_model( 301 QuantizationDebuggerTest.tf_model_root, 302 QuantizationDebuggerTest.tf_model, 303 _calibration_gen, 304 debug=False) 305 306 with self.assertRaisesRegex( 307 ValueError, 'Please check if the quantized model is in debug mode'): 308 debugger.QuantizationDebugger( 309 quant_debug_model_content=normal_quant_model, 310 debug_dataset=_calibration_gen) 311 312 @parameterized.named_parameters( 313 ('empty quantization parameter', { 314 'quantization_parameters': {} 315 }, None), 316 ('empty scales/zero points', { 317 'quantization_parameters': { 318 'scales': [], 319 'zero_points': [] 320 } 321 }, None), 322 ('invalid scales/zero points', { 323 'quantization_parameters': { 324 'scales': [1.0], 325 'zero_points': [] 326 } 327 }, None), 328 ('correct case', { 329 'quantization_parameters': { 330 'scales': [0.5, 1.0], 331 'zero_points': [42, 7] 332 } 333 }, (0.5, 42)), 334 ) 335 def test_get_quant_params(self, tensor_detail, expected_value): 336 self.assertEqual(debugger._get_quant_params(tensor_detail), expected_value) 337 338 @parameterized.named_parameters( 339 ('float_io', False), 340 ('quantized_io', True)) 341 @test_util.run_v2_only 342 def test_denylisted_ops_from_option_setter(self, quantized_io): 343 options = debugger.QuantizationDebugOptions( 344 layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))}, 345 fully_quantize=quantized_io) 346 quant_debugger = debugger.QuantizationDebugger( 347 converter=_quantize_converter(self.tf_model_root, self.tf_model, 348 _calibration_gen), 349 debug_dataset=_calibration_gen, 350 debug_options=options) 351 352 options.denylisted_ops = ['CONV_2D'] 353 # TODO(b/195084873): The exception is expected to check whether selective 354 # quantization was done properly, since after the selective quantization 355 # the model will have no quantized layers thus have no NumericVerify ops, 356 # resulted in this exception. Marked with a bug to fix this in more 357 # straightforward way. 358 with self.assertRaisesRegex( 359 ValueError, 'Please check if the quantized model is in debug mode'): 360 quant_debugger.options = options 361 362 @parameterized.named_parameters( 363 ('float_io', False), 364 ('quantized_io', True)) 365 @test_util.run_v2_only 366 def test_denylisted_ops_from_option_constructor(self, quantized_io): 367 options = debugger.QuantizationDebugOptions( 368 layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))}, 369 fully_quantize=quantized_io, 370 denylisted_ops=['CONV_2D']) 371 # TODO(b/195084873): Count the number of NumericVerify op. 372 with self.assertRaisesRegex( 373 ValueError, 'Please check if the quantized model is in debug mode'): 374 _ = debugger.QuantizationDebugger( 375 converter=_quantize_converter(self.tf_model_root, self.tf_model, 376 _calibration_gen), 377 debug_dataset=_calibration_gen, 378 debug_options=options) 379 380 @parameterized.named_parameters(('float_io', False), ('quantized_io', True)) 381 @test_util.run_v2_only 382 def test_denylisted_nodes_from_option_setter(self, quantized_io): 383 options = debugger.QuantizationDebugOptions( 384 layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))}, 385 fully_quantize=quantized_io) 386 quant_debugger = debugger.QuantizationDebugger( 387 converter=_quantize_converter(self.tf_model_root, self.tf_model, 388 _calibration_gen), 389 debug_dataset=_calibration_gen, 390 debug_options=options) 391 392 options.denylisted_nodes = ['Identity'] 393 # TODO(b/195084873): Count the number of NumericVerify op. 394 with self.assertRaisesRegex( 395 ValueError, 'Please check if the quantized model is in debug mode'): 396 quant_debugger.options = options 397 398 @parameterized.named_parameters(('float_io', False), ('quantized_io', True)) 399 @test_util.run_v2_only 400 def test_denylisted_nodes_from_option_constructor(self, quantized_io): 401 options = debugger.QuantizationDebugOptions( 402 layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))}, 403 fully_quantize=quantized_io, 404 denylisted_nodes=['Identity']) 405 # TODO(b/195084873): Count the number of NumericVerify op. 406 with self.assertRaisesRegex( 407 ValueError, 'Please check if the quantized model is in debug mode'): 408 _ = debugger.QuantizationDebugger( 409 converter=_quantize_converter(self.tf_model_root, self.tf_model, 410 _calibration_gen), 411 debug_dataset=_calibration_gen, 412 debug_options=options) 413 414 @mock.patch.object(metrics.TFLiteMetrics, 415 'increase_counter_debugger_creation') 416 def test_creation_counter(self, increase_call): 417 debug_model = QuantizationDebuggerTest.debug_model_float 418 debugger.QuantizationDebugger( 419 quant_debug_model_content=debug_model, debug_dataset=_calibration_gen) 420 increase_call.assert_called_once() 421 422 423if __name__ == '__main__': 424 test.main() 425